diff --git a/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_10.txt b/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_10.txt
index e8238d6e0c9d98..91abdd8adf1c86 100644
--- a/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_10.txt
+++ b/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_10.txt
@@ -248,9 +248,9 @@ libclang==18.1.1 \
 lit==17.0.6 \
     --hash=sha256:dfa9af9b55fc4509a56be7bf2346f079d7f4a242d583b9f2e0b078fd0abae31b
     # via -r ci/official/requirements_updater/requirements.in
-markdown==3.8 \
-    --hash=sha256:794a929b79c5af141ef5ab0f2f642d0f7b1872981250230e72682346f7cc90dc \
-    --hash=sha256:7df81e63f0df5c4b24b7d156eb81e4690595239b7d70937d0409f1b0de319c6f
+markdown==3.8.1 \
+    --hash=sha256:46cc0c0f1e5211ab2e9d453582f0b28a1bfaf058a9f7d5c50386b99b588d8811 \
+    --hash=sha256:a2e2f01cead4828ee74ecca9623045f62216aef2212a7685d6eb9163f590b8c1
     # via tb-nightly
 markdown-it-py==3.0.0 \
     --hash=sha256:355216845c60bd96232cd8d8c40e8f9765cc86f46880e43a8fd22dc1a1a8cab1 \
diff --git a/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_11.txt b/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_11.txt
index 9237c82440f566..7a07aa010ad0ad 100644
--- a/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_11.txt
+++ b/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_11.txt
@@ -248,9 +248,9 @@ libclang==18.1.1 \
 lit==17.0.6 \
     --hash=sha256:dfa9af9b55fc4509a56be7bf2346f079d7f4a242d583b9f2e0b078fd0abae31b
     # via -r ci/official/requirements_updater/requirements.in
-markdown==3.8 \
-    --hash=sha256:794a929b79c5af141ef5ab0f2f642d0f7b1872981250230e72682346f7cc90dc \
-    --hash=sha256:7df81e63f0df5c4b24b7d156eb81e4690595239b7d70937d0409f1b0de319c6f
+markdown==3.8.1 \
+    --hash=sha256:46cc0c0f1e5211ab2e9d453582f0b28a1bfaf058a9f7d5c50386b99b588d8811 \
+    --hash=sha256:a2e2f01cead4828ee74ecca9623045f62216aef2212a7685d6eb9163f590b8c1
     # via tb-nightly
 markdown-it-py==3.0.0 \
     --hash=sha256:355216845c60bd96232cd8d8c40e8f9765cc86f46880e43a8fd22dc1a1a8cab1 \
diff --git a/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_12.txt b/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_12.txt
index 12559bc05817da..d3cb6403cb4979 100644
--- a/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_12.txt
+++ b/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_12.txt
@@ -248,9 +248,9 @@ libclang==18.1.1 \
 lit==17.0.6 \
     --hash=sha256:dfa9af9b55fc4509a56be7bf2346f079d7f4a242d583b9f2e0b078fd0abae31b
     # via -r ci/official/requirements_updater/requirements.in
-markdown==3.8 \
-    --hash=sha256:794a929b79c5af141ef5ab0f2f642d0f7b1872981250230e72682346f7cc90dc \
-    --hash=sha256:7df81e63f0df5c4b24b7d156eb81e4690595239b7d70937d0409f1b0de319c6f
+markdown==3.8.1 \
+    --hash=sha256:46cc0c0f1e5211ab2e9d453582f0b28a1bfaf058a9f7d5c50386b99b588d8811 \
+    --hash=sha256:a2e2f01cead4828ee74ecca9623045f62216aef2212a7685d6eb9163f590b8c1
     # via tb-nightly
 markdown-it-py==3.0.0 \
     --hash=sha256:355216845c60bd96232cd8d8c40e8f9765cc86f46880e43a8fd22dc1a1a8cab1 \
diff --git a/tensorflow/core/data/service/server_lib.cc b/tensorflow/core/data/service/server_lib.cc
index e0b5f3cf9ac90b..1555fe0c9ac012 100644
--- a/tensorflow/core/data/service/server_lib.cc
+++ b/tensorflow/core/data/service/server_lib.cc
@@ -103,7 +103,12 @@ void GrpcDataServerBase::Stop() {
   stopped_ = true;
 }
 
-void GrpcDataServerBase::Join() { server_->Wait(); }
+void GrpcDataServerBase::Join() {
+  if (!server_) {
+    return;
+  }
+  server_->Wait();
+}
 
 int GrpcDataServerBase::BoundPort() { return bound_port(); }
 
diff --git a/tensorflow/core/data/service/server_lib.h b/tensorflow/core/data/service/server_lib.h
index 56a8f8d94fc558..8647da848591ea 100644
--- a/tensorflow/core/data/service/server_lib.h
+++ b/tensorflow/core/data/service/server_lib.h
@@ -124,7 +124,7 @@ class DispatchGrpcDataServer : public GrpcDataServerBase {
  private:
   const experimental::DispatcherConfig config_;
   // Owned. We use a raw pointer because GrpcDispatcherImpl is forward-declared.
-  GrpcDispatcherImpl* service_;
+  GrpcDispatcherImpl* service_ = nullptr;
 };
 
 // A wrapper for `SnapshotTaskProgress` for use with pybind.
@@ -170,7 +170,7 @@ class WorkerGrpcDataServer : public GrpcDataServerBase {
 
   const experimental::WorkerConfig config_;
   // Owned. We use a raw pointer because GrpcWorkerImpl is forward-declared.
-  GrpcWorkerImpl* service_;
+  GrpcWorkerImpl* service_ = nullptr;
   std::shared_ptr<DataTransferServer> transfer_server_;
 };
 
diff --git a/tensorflow/core/kernels/linalg/BUILD b/tensorflow/core/kernels/linalg/BUILD
index 31a192ee9a5c12..4e16037f225206 100644
--- a/tensorflow/core/kernels/linalg/BUILD
+++ b/tensorflow/core/kernels/linalg/BUILD
@@ -134,7 +134,10 @@ tf_kernel_library(
 tf_kernel_library(
     name = "matrix_diag_op",
     prefix = "matrix_diag_op",
-    deps = LINALG_DEPS,
+    deps = LINALG_DEPS + [
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+    ],
 )
 
 tf_kernel_library(
@@ -183,7 +186,10 @@ tf_kernel_library(
 tf_kernel_library(
     name = "self_adjoint_eig_op",
     prefix = "self_adjoint_eig_op",
-    deps = LINALG_DEPS + ["//tensorflow/core:lib_internal"],
+    deps = LINALG_DEPS + [
+        "//tensorflow/core:lib_internal",
+        "@com_google_absl//absl/status",
+    ],
 )
 
 tf_kernel_library(
diff --git a/tensorflow/core/kernels/linalg/lu_op.cc b/tensorflow/core/kernels/linalg/lu_op.cc
index 998769e31fb0e0..de0f0d2ad1a322 100644
--- a/tensorflow/core/kernels/linalg/lu_op.cc
+++ b/tensorflow/core/kernels/linalg/lu_op.cc
@@ -17,6 +17,8 @@ limitations under the License.
 #include <limits>
 
 #include "absl/container/inlined_vector.h"
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
 #include "Eigen/Core"  // from @eigen_archive
 #include "Eigen/LU"  // from @eigen_archive
 #include "tensorflow/core/framework/kernel_def_builder.h"
diff --git a/tensorflow/core/kernels/linalg/matrix_diag_op.cc b/tensorflow/core/kernels/linalg/matrix_diag_op.cc
index 90fa24b8122ef7..f017da44bd9208 100644
--- a/tensorflow/core/kernels/linalg/matrix_diag_op.cc
+++ b/tensorflow/core/kernels/linalg/matrix_diag_op.cc
@@ -14,17 +14,19 @@ limitations under the License.
 ==============================================================================*/
 
 // See docs in ../ops/array_ops.cc.
+#include <string>
+#include <tuple>
+#include <utility>
+
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
 #define EIGEN_USE_THREADS
 
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
 #define EIGEN_USE_GPU
 #endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
 
-#include "tensorflow/core/kernels/linalg/matrix_diag_op.h"
-
 #include <algorithm>
-#include <memory>
-#include <vector>
 
 #include "unsupported/Eigen/CXX11/Tensor"  // from @eigen_archive
 #include "tensorflow/core/framework/op_kernel.h"
@@ -33,6 +35,7 @@ limitations under the License.
 #include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/framework/tensor_types.h"
 #include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/kernels/linalg/matrix_diag_op.h"
 #include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/lib/core/threadpool.h"
 #include "tensorflow/core/platform/logging.h"
diff --git a/tensorflow/core/kernels/linalg/self_adjoint_eig_op.cc b/tensorflow/core/kernels/linalg/self_adjoint_eig_op.cc
index 816fb97174831f..478d3174e2fca2 100644
--- a/tensorflow/core/kernels/linalg/self_adjoint_eig_op.cc
+++ b/tensorflow/core/kernels/linalg/self_adjoint_eig_op.cc
@@ -17,6 +17,7 @@ limitations under the License.
 
 #include <cstdint>
 
+#include "absl/status/status.h"
 #include "Eigen/Core"  // from @eigen_archive
 #include "Eigen/Eigenvalues"  // from @eigen_archive
 #include "tensorflow/core/framework/kernel_def_builder.h"
diff --git a/tensorflow/python/data/experimental/service/server_lib_test.py b/tensorflow/python/data/experimental/service/server_lib_test.py
index 6d44747def1668..2866a549bc9008 100644
--- a/tensorflow/python/data/experimental/service/server_lib_test.py
+++ b/tensorflow/python/data/experimental/service/server_lib_test.py
@@ -148,6 +148,28 @@ def testJoinWorker(self):
     worker.stop()
     worker.join()
 
+  def testJoinDispatcherWithoutStart(self):
+    dispatcher = server_lib.DispatchServer(start=False)
+    dispatcher.join()
+
+  def testJoinWorkerWithoutStart(self):
+    dispatcher = server_lib.DispatchServer()
+    worker = server_lib.WorkerServer(
+        server_lib.WorkerConfig(dispatcher._address), start=False
+    )
+    worker.join()
+
+  def testDispatcherDestroyedWithoutStart(self):
+    dispatcher = server_lib.DispatchServer(start=False)
+    del dispatcher
+
+  def testWorkerDestroyedWithoutStart(self):
+    dispatcher = server_lib.DispatchServer()
+    worker = server_lib.WorkerServer(
+        server_lib.WorkerConfig(dispatcher._address), start=False
+    )
+    del worker
+
   def testDispatcherNumWorkers(self):
     dispatcher = server_lib.DispatchServer()
     self.assertEqual(0, dispatcher._num_workers())
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/BUILD b/third_party/xla/xla/backends/gpu/codegen/triton/BUILD
index 34cf8bfd3ca081..ec3ee6dfe836c5 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/BUILD
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/BUILD
@@ -366,7 +366,6 @@ xla_test(
         "//xla/backends/gpu/tests:gpu_pjrt_codegen_test",
         "//xla/backends/gpu/transforms:convert_triton_gemm_config",
         "//xla/codegen/xtile:block_level_parameters",
-        "//xla/hlo/analysis:symbolic_map",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/testlib:filecheck",
         "//xla/hlo/testlib:pattern_matcher_gmock",
@@ -379,8 +378,10 @@ xla_test(
         "//xla/service/gpu:target_constants",
         "//xla/stream_executor:device_description",
         "//xla/stream_executor/cuda:cuda_compute_capability",
-        "//xla/tests:hlo_pjrt_interpreter_reference_mixin",
+        "//xla/tests:hlo_interpreter_reference_mixin",
         "//xla/tests:xla_internal_test_main",  # fixdeps: keep
+        "//xla/tsl/lib/core:status_test_util",
+        "//xla/tsl/platform:env",
         "//xla/tsl/platform:errors",
         "//xla/tsl/platform:test",
         "@com_google_absl//absl/log:check",
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/triton_gemm_fusion_test.cc b/third_party/xla/xla/backends/gpu/codegen/triton/triton_gemm_fusion_test.cc
index 5ffa216b50edda..3186968ae3c3b3 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/triton_gemm_fusion_test.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/triton_gemm_fusion_test.cc
@@ -40,7 +40,6 @@ limitations under the License.
 #include "xla/backends/gpu/transforms/convert_triton_gemm_config.h"
 #include "xla/codegen/xtile/block_level_parameters.h"
 #include "xla/error_spec.h"
-#include "xla/hlo/analysis/symbolic_expr.h"
 #include "xla/hlo/ir/hlo_casting_utils.h"
 #include "xla/hlo/ir/hlo_computation.h"
 #include "xla/hlo/ir/hlo_instruction.h"
@@ -57,7 +56,9 @@ limitations under the License.
 #include "xla/service/pattern_matcher.h"
 #include "xla/stream_executor/cuda/cuda_compute_capability.h"
 #include "xla/stream_executor/device_description.h"
-#include "xla/tests/hlo_pjrt_interpreter_reference_mixin.h"
+#include "xla/tests/hlo_interpreter_reference_mixin.h"
+#include "xla/tsl/lib/core/status_test_util.h"
+#include "xla/tsl/platform/env.h"
 #include "xla/tsl/platform/errors.h"
 #include "xla/tsl/platform/test.h"
 #include "xla/xla.pb.h"
@@ -87,9 +88,9 @@ HloInstruction* GetNonBitcastRoot(const HloComputation* computation) {
   return root;
 }
 
-class TritonTest : public HloInterpreterReferenceMixin<GpuPjRtCodegenTest> {
+class TritonTestBase : public HloInterpreterReferenceMixin<GpuPjRtCodegenTest> {
  public:
-  TritonTest() = default;
+  TritonTestBase() = default;
   DebugOptions GetDebugOptionsForTest() const override {
     DebugOptions debug_options = GpuPjRtCodegenTest::GetDebugOptionsForTest();
     debug_options.set_xla_gpu_autotune_level(0);
@@ -134,6 +135,16 @@ class TritonTest : public HloInterpreterReferenceMixin<GpuPjRtCodegenTest> {
   mlir::MLIRContext mlir_context_;
 };
 
+class TritonTest : public TritonTestBase,
+                   public ::testing::WithParamInterface<bool> {
+ public:
+  DebugOptions GetDebugOptionsForTest() const override {
+    DebugOptions debug_options = TritonTestBase::GetDebugOptionsForTest();
+    debug_options.set_xla_gpu_experimental_gemm_fusion_v2(GetParam());
+    return debug_options;
+  }
+};
+
 class TritonGemmTest : public TritonTest {
  public:
   DebugOptions GetDebugOptionsForTest() const override {
@@ -158,7 +169,16 @@ class TritonGemmTest : public TritonTest {
   }
 };
 
-TEST_F(TritonTest, TestGemmWithTrivialNonContractingDimension) {
+std::string GemmFusionVersionName(const ::testing::TestParamInfo<bool>& info) {
+  return info.param ? "GemmFusionV2" : "GemmFusionV1";
+}
+
+INSTANTIATE_TEST_SUITE_P(All, TritonTest, ::testing::Bool(),
+                         GemmFusionVersionName);
+INSTANTIATE_TEST_SUITE_P(All, TritonGemmTest, ::testing::Bool(),
+                         GemmFusionVersionName);
+
+TEST_P(TritonTest, TestGemmWithTrivialNonContractingDimension) {
   constexpr absl::string_view kHloText = R"(
 HloModule t, is_scheduled=true
 
@@ -189,7 +209,7 @@ CHECK: tt.dot {{.*}} : tensor<16x32xf32> * tensor<32x16xf32> -> tensor<16x16xf32
 )"));
 }
 
-TEST_F(TritonTest, PredParametersAreTruncatedToI1) {
+TEST_P(TritonTest, PredParametersAreTruncatedToI1) {
   constexpr absl::string_view kHloText = R"(
 HloModule m
 
@@ -233,7 +253,7 @@ CHECK: %{{.*}} = arith.andi %[[CMPI]], %{{.*}} : tensor<16x16xi1>
 )"));
 }
 
-TEST_F(TritonTest, CodegenBatchedDotWithConcatenationWithCorrectBatchStride) {
+TEST_P(TritonTest, CodegenBatchedDotWithConcatenationWithCorrectBatchStride) {
   constexpr absl::string_view kHloText = R"(
 HloModule t
 
@@ -270,7 +290,7 @@ CHECK: tt.dot {{.*}} : tensor<16x32xf32> * tensor<32x64xf32> -> tensor<16x64xf32
 }
 
 // TODO(b/417172838): enable after enabling dynamic slice in support.cc.
-TEST_F(TritonTest, DISABLED_CodegenDynamicSliceWithCorrectOffsets) {
+TEST_P(TritonTest, DISABLED_CodegenDynamicSliceWithCorrectOffsets) {
   // TODO(b/417172838): we now should support non-majormost dimensions, port
   // this test to fusion_emitter_device_test with that support.
 
@@ -325,7 +345,7 @@ ENTRY e {
   )"));
 }
 
-TEST_F(TritonGemmTest, DoNotUseTensorCoresWithHighestPrecision) {
+TEST_P(TritonGemmTest, DoNotUseTensorCoresWithHighestPrecision) {
   constexpr absl::string_view kHloText = R"(
 triton_gemm_r {
   parameter_0 = s8[80,15]{1,0} parameter(0)
@@ -355,7 +375,7 @@ CHECK-NOT: mma
 )");
 }
 
-TEST_F(TritonGemmTest, UseTensorCoresWithHighPrecision) {
+TEST_P(TritonGemmTest, UseTensorCoresWithHighPrecision) {
   constexpr absl::string_view kHloText = R"(
 triton_gemm_r {
   parameter_0 = s8[80,15]{1,0} parameter(0)
@@ -385,7 +405,7 @@ CHECK: mma
 )");
 }
 
-TEST_F(TritonGemmTest, DebugOptionsArePropagated) {
+TEST_P(TritonGemmTest, DebugOptionsArePropagated) {
   constexpr absl::string_view kHloText = R"(
 ENTRY e {
   p0 = f16[30,30] parameter(0)
@@ -423,7 +443,7 @@ ENTRY e {
   EXPECT_GT(file_size, 10);
 }
 
-TEST_F(TritonGemmTest, DotWithPredFromCompareProducesCorrectResult) {
+TEST_P(TritonGemmTest, DotWithPredFromCompareProducesCorrectResult) {
   constexpr absl::string_view kHloText = R"(
 triton_dot {
   parameter_0 = s32[4,128]{1,0} parameter(0)
@@ -449,7 +469,7 @@ ENTRY main {
                                ErrorSpec{/*aabs=*/1e-4, /*arel=*/1e-6}));
 }
 
-TEST_F(TritonGemmTest, UseTensorCoresForF32OnAmpere) {
+TEST_P(TritonGemmTest, UseTensorCoresForF32OnAmpere) {
   constexpr absl::string_view kHloText = R"(
 triton_gemm_r {
   parameter_0 = f16[80,15]{1,0} parameter(0)
@@ -479,7 +499,7 @@ CHECK: mma
 
 // TODO(b/353484968): Tests that don't run RunAndCompareNoHloPasses should be
 // moved to deviceless test file.
-TEST_F(TritonGemmTest, FailIfTooMuchShmem) {
+TEST_P(TritonGemmTest, FailIfTooMuchShmem) {
   auto cc = se::CudaComputeCapability::Ampere();
   const se::DeviceDescription device_info =
       TestGpuDeviceInfo::RTXA6000DeviceInfo();
@@ -536,7 +556,7 @@ ENTRY entry {
   EXPECT_GT(result.shmem_bytes, device_info.shared_memory_per_block());
 }
 
-TEST_F(TritonGemmTest, MultipleDims) {
+TEST_P(TritonGemmTest, MultipleDims) {
   constexpr absl::string_view kHloText = R"(
 HloModule t
 
@@ -559,7 +579,7 @@ ENTRY e {
   EXPECT_TRUE(RunAndCompare(kHloText, ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3}));
 }
 
-TEST_F(TritonGemmTest, MultipleBatchDimensions) {
+TEST_P(TritonGemmTest, MultipleBatchDimensions) {
   constexpr absl::string_view kHloText = R"(
 HloModule m
 
@@ -593,7 +613,7 @@ ENTRY e {
       std::move(module), ErrorSpec{/*aabs=*/2e-2, /*arel=*/2e-2}));
 }
 
-TEST_F(TritonGemmTest, PredWithBF16DotProducesCorrectResult) {
+TEST_P(TritonGemmTest, PredWithBF16DotProducesCorrectResult) {
   constexpr absl::string_view kHloText = R"(
 triton_dot {
   p0 = pred[8,640]{1,0} parameter(0)
@@ -622,7 +642,7 @@ ENTRY e {
 // TODO: b/422676780 - Enable the tests once the indexing maps-based tiling is
 // deprecated. The test is disabled after we remove TransposeDimensionGrouper
 // pass, because the infra currently requires grouping of adjacent dimensions.
-TEST_F(TritonGemmTest, DISABLED_SplitLhsNoncontractingTransposeRhs) {
+TEST_P(TritonGemmTest, DISABLED_SplitLhsNoncontractingTransposeRhs) {
   constexpr absl::string_view kHloText = R"(
 HloModule t
 
@@ -650,7 +670,7 @@ ENTRY e {
 // TODO: b/422676780 - Enable the tests once the indexing maps-based tiling is
 // deprecated. The test is disabled after we remove TransposeDimensionGrouper
 // pass, because the infra currently requires grouping of adjacent dimensions.
-TEST_F(TritonGemmTest, DISABLED_SplitLhsNoncontracting) {
+TEST_P(TritonGemmTest, DISABLED_SplitLhsNoncontracting) {
   constexpr absl::string_view kHloText = R"(
 ENTRY e {
   p0 = f32[72,72] parameter(0)
@@ -677,7 +697,7 @@ ENTRY e {
   EXPECT_TRUE(RunAndCompare(kHloText, ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3}));
 }
 
-TEST_F(TritonGemmTest, SplitAndTransposeLhsExecutesCorrectly) {
+TEST_P(TritonGemmTest, SplitAndTransposeLhsExecutesCorrectly) {
   constexpr absl::string_view kHloText = R"(
 HloModule m
 
@@ -710,7 +730,7 @@ ENTRY e {
 // TODO(b/393299275): it's not clear that this test is actually testing what it
 // claims to be testing. It should either be rewritten to start from
 // post-optimization HLO, or hoisted out to test the fusion logic specifically.
-TEST_F(TritonGemmTest, NondefaultOperandLayoutIsSupported) {
+TEST_P(TritonGemmTest, NondefaultOperandLayoutIsSupported) {
   constexpr absl::string_view kHloText = R"(
 ENTRY r {
   p1 = f16[3,10,128]{2,1,0} parameter(1)
@@ -723,10 +743,10 @@ ENTRY r {
 })";
 
   MatchOptimizedHlo(kHloText, R"(
-; CHECK: %[[p0:.*]] = f16[10,3,128]{2,0,1} parameter(0)
+; CHECK: %[[p0:.*]] = f16[10,3,128]{2,0,1} parameter
 ; CHECK: %[[cv:.*]] = f32[10,3,128]{2,0,1} convert(%[[p0]])
 ; CHECK: %[[tr:.*]] = f32[3,10,128]{2,0,1} transpose(%[[cv]]), dimensions={1,0,2}
-; CHECK: %[[p1:.*]] = f32[3,10,123]{2,1,0} parameter(1)
+; CHECK: %[[p1:.*]] = f32[3,10,123]{2,1,0} parameter
 ; CHECK: f32[3,128,123]{2,1,0} dot(%[[tr]], %[[p1]])
 )");
 
@@ -735,7 +755,7 @@ ENTRY r {
 
 // TODO(b/393299275): this is a pure test of fusion logic. It should be moved to
 // a separate, fusion-specific, deviceless test.
-TEST_F(TritonGemmTest, DoNotFuseSplitRhsContractingTranspose) {
+TEST_P(TritonGemmTest, DoNotFuseSplitRhsContractingTranspose) {
   constexpr absl::string_view kHloText = R"(
 HloModule t
 
@@ -761,7 +781,7 @@ ENTRY e {
 
 // TODO(b/393299275): this is a pure test of fusion logic. It should be moved to
 // a separate, fusion-specific, deviceless test.
-TEST_F(TritonGemmTest, DoNotFuseSplitLhsContractingTranspose) {
+TEST_P(TritonGemmTest, DoNotFuseSplitLhsContractingTranspose) {
   constexpr absl::string_view kHloText = R"(
 HloModule t
 
@@ -788,7 +808,7 @@ ENTRY e {
 // TODO(b/393299275): this test should be rewritten to start from
 // post-optimization HLO. (Though I'm not entirely sure it's even worth keeping
 // it.)
-TEST_F(TritonGemmTest, NonMajorMostInputBatchWorksCorrectly) {
+TEST_P(TritonGemmTest, NonMajorMostInputBatchWorksCorrectly) {
   constexpr absl::string_view kHloText = R"(
 HloModule t
 
@@ -815,7 +835,7 @@ ENTRY e {
 
 // TODO(b/393299275): this is a pure test of fusion logic. It should be moved to
 // a separate, fusion-specific, deviceless test.
-TEST_F(TritonGemmTest, DoNotFuseArbitraryReshape) {
+TEST_P(TritonGemmTest, DoNotFuseArbitraryReshape) {
   constexpr absl::string_view kHloText = R"(
 HloModule m
 
@@ -838,36 +858,12 @@ ENTRY e {
 )");
 }
 
-TEST_F(TritonGemmTest, MultipleBatchRequireSeparateTranspose) {
-  constexpr absl::string_view kHloText = R"(
-HloModule m
-
-ENTRY e {
-  Arg_0 = f16[3,4,2,5,4] parameter(0)
-  c = f32[3,4,2,5,4] convert(Arg_0)
-  Arg_1 = f32[5,3,4,3,2] parameter(1)
-  ROOT dot.3 = f32[5,3,4,4,3] dot(c, Arg_1),
-    lhs_batch_dims={3,0,1}, lhs_contracting_dims={2},
-    rhs_batch_dims={0,1,2}, rhs_contracting_dims={4}
-})";
-
-  MatchOptimizedHlo(kHloText, R"(
-; CHECK: ROOT
-; CHECK: transpose(
-; CHECK: bitcast(
-; CHECK: kCustom
-; CHECK-SAME: backend_config={{.*}}"kind":"__triton_nested_gemm_fusion"
-)");
-
-  EXPECT_TRUE(RunAndCompare(kHloText, ErrorSpec{/*aabs=*/1e-4, /*arel=*/1e-4}));
-}
-
 // This tests the complexity heuristics in TritonWrapper.
 // TODO(b/393299275): this is not worth keeping as a codegen test. Really, we
 // should not reject tilings that are slow/spill in codegen. If this has use in
 // autotuning, then this should be tested/called in the autotuner.
 // The generic Triton emitter does not want to deal with this.
-TEST_F(TritonGemmTest, DISABLED_FailForTooComplexTiling) {
+TEST_P(TritonGemmTest, DISABLED_FailForTooComplexTiling) {
   auto cc = se::CudaComputeCapability::Ampere();
   const se::DeviceDescription device_info =
       TestGpuDeviceInfo::RTXA6000DeviceInfo();
@@ -928,7 +924,7 @@ ENTRY entry {
 // TODO(b/393299275): this test may have some value while Triton tiling
 // propagation is being replaced, but has little worth as a codegen test.
 // Consider moving this.
-TEST_F(TritonGemmTest,
+TEST_P(TritonGemmTest,
        BroadcastsOfTriviallySizedNonContractingDimensionsAreSupported) {
   constexpr absl::string_view kHloText = R"(
 f {
@@ -963,7 +959,7 @@ e {
 // TODO(b/393299275): this test may have some value while Triton tiling
 // propagation is being replaced, but has little worth as a codegen test.
 // Consider moving this.
-TEST_F(TritonGemmTest,
+TEST_P(TritonGemmTest,
        BroadcastsOfTriviallySizedContractingDimensionsAreSupported) {
   constexpr absl::string_view kHloText = R"(
 f {
@@ -994,7 +990,7 @@ e {
 }
 
 // TODO(b/417172838): enable after enabling dynamic slice in support.cc.
-TEST_F(TritonGemmTest, DISABLED_DynamicSliceIsSupportedInLhsEndToEnd) {
+TEST_P(TritonGemmTest, DISABLED_DynamicSliceIsSupportedInLhsEndToEnd) {
   // The select is used to restrict the start index to values that make sense.
   // If it was constant, then the dynamic-slice would be optimized to slice. It
   // is not strictly needed, because we also support clamping the indices.
@@ -1029,7 +1025,7 @@ ENTRY e {
 }
 
 // TODO(b/417172838): enable after enabling dynamic slice in support.cc.
-TEST_F(TritonGemmTest, DISABLED_DynamicSliceIsSupportedInRhs) {
+TEST_P(TritonGemmTest, DISABLED_DynamicSliceIsSupportedInRhs) {
   // The start index(es) for the non-majormost dimension(s) are constant zero(s)
   // because we don't support dynamic slice on those dimensions.
   constexpr absl::string_view kHloText = R"(
@@ -1067,7 +1063,7 @@ ENTRY e {
 }
 
 class TritonGemmDynamicSliceClampingTest
-    : public TritonTest,
+    : public TritonTestBase,
       public ::testing::WithParamInterface<int> {};
 
 // TODO(b/417172838): enable after enabling dynamic slice in support.cc.
@@ -1122,7 +1118,7 @@ INSTANTIATE_TEST_SUITE_P(All, TritonGemmDynamicSliceClampingTest,
                          ::testing::Values(-100, 3, 999), OffsetParamToString);
 
 // TODO(b/417172838): enable after enabling dynamic slice in support.cc.
-TEST_F(TritonGemmTest,
+TEST_P(TritonGemmTest,
        DISABLED_DynamicSliceOfMajormostContractingDimIsSupported) {
   // Tests that dynamic-slice works on the majormost dimension even if that
   // dimension is contracted.
@@ -1163,7 +1159,7 @@ ENTRY e {
 }
 
 // TODO(b/417172838): enable after enabling dynamic slice in support.cc.
-TEST_F(TritonGemmTest, DISABLED_DynamicSliceOfMajormostBatchDimIsSupported) {
+TEST_P(TritonGemmTest, DISABLED_DynamicSliceOfMajormostBatchDimIsSupported) {
   // Tests that dynamic-slice works on the majormost dimension even if that
   // dimension is a batch.
   // The start index(es) for the non-majormost dimension(s) are constant zero(s)
@@ -1208,7 +1204,7 @@ ENTRY e {
 }
 
 // TODO(b/417172838): enable after enabling dynamic slice in support.cc.
-TEST_F(TritonGemmTest,
+TEST_P(TritonGemmTest,
        DISABLED_DynamicSliceSingleDimensionIntoReshapeIsSupported) {
   // This directly tests the targeted use case (b/307922364) of iterating over
   // layer weights and extracting them with dynamic slice.
@@ -1253,7 +1249,7 @@ ENTRY e {
 
 // TODO(b/393299275): this should just be a fusion test and does not need to be
 // in the codegen directory.
-TEST_F(TritonGemmTest, DoNotFuseConcatenationOfSplitNonContractingDimension) {
+TEST_P(TritonGemmTest, DoNotFuseConcatenationOfSplitNonContractingDimension) {
   if (GpuComputeCapability().IsRocm()) {
     GTEST_SKIP() << "Not using autotuner on ROCM yet.";
   }
@@ -1283,7 +1279,7 @@ ENTRY e {
   EXPECT_TRUE(RunAndCompare(kHloText, ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3}));
 }
 
-TEST_F(TritonGemmTest, BinaryOperationWithSmallInputsIsFused) {
+TEST_P(TritonGemmTest, BinaryOperationWithSmallInputsIsFused) {
   constexpr absl::string_view kHloText = R"(
 HloModule m
 
@@ -1307,7 +1303,7 @@ ENTRY e {
                      .WithFusionKind(HloInstruction::FusionKind::kCustom)));
 }
 
-TEST_F(TritonGemmTest, BinaryOperationWithLargeInputsIsNotFused) {
+TEST_P(TritonGemmTest, BinaryOperationWithLargeInputsIsNotFused) {
   constexpr absl::string_view kHloText = R"(
 HloModule m
 
@@ -1336,7 +1332,7 @@ ENTRY e {
 )");
 }
 
-TEST_F(TritonGemmTest, ParametersWithDifferentLayoutsAreSupportedInOneScope) {
+TEST_P(TritonGemmTest, ParametersWithDifferentLayoutsAreSupportedInOneScope) {
   constexpr absl::string_view kHloText = R"(
 ENTRY e {
   p0 = s8[5,3] parameter(0)
@@ -1357,7 +1353,7 @@ ENTRY e {
                      .WithFusionKind(HloInstruction::FusionKind::kCustom)));
 }
 
-TEST_F(TritonGemmTest, BinaryOperationOnLargeParametersIsFused) {
+TEST_P(TritonGemmTest, BinaryOperationOnLargeParametersIsFused) {
   constexpr absl::string_view kHloText = R"(
 HloModule m
 
@@ -1380,45 +1376,55 @@ ENTRY e {
                      .WithFusionKind(HloInstruction::FusionKind::kCustom)));
 }
 
-TEST_F(TritonGemmTest, LinkingLibdeviceTwiceWorks) {
-  // TODO(b/393299275): This test looks weird. It's testing the whole
-  // optimization pipeline end-to-end to check that linking libdevice twice
-  // works? rewrite this to just use post-optimization HLO
+TEST_P(TritonGemmTest, LinkingLibdeviceTwiceWorks) {
   constexpr absl::string_view kHloText = R"(
-ENTRY e {
-  p0 = s8[7,3] parameter(0)
-  c0 = f32[7,3] convert(p0)
-  p1 = f32[3,16] parameter(1)
-  e1 = f32[3,16] exponential(p1)
-  d0 = f32[7,16] dot(c0, e1),
-    lhs_contracting_dims={1}, rhs_contracting_dims={0}
-  p2 = s8[7,3] parameter(2)
-  c2 = f32[7,3] convert(p2)
-  e2 = f32[7,3] exponential(c2)
-  p3 = f32[3,16] parameter(3)
-  d1 = f32[7,16] dot(e2, p3),
+HloModule m
+
+triton_dot_0 {
+  p0 = s8[7,3]{1,0} parameter(0)
+  c0 = f32[7,3]{1,0} convert(p0)
+  p1 = f32[3,16]{1,0} parameter(1)
+  e1 = f32[3,16]{1,0} exponential(p1)
+  ROOT d0 = f32[7,16]{1,0} dot(c0, e1),
     lhs_contracting_dims={1}, rhs_contracting_dims={0}
-  ROOT a = f32[7,16] add(d0, d1)
-})";
+}
 
-  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
-                       GetOptimizedModule(kHloText));
+triton_dot_1 {
+  p0 = s8[7,3]{1,0} parameter(0)
+  c0 = f32[7,3]{1,0} convert(p0)
+  e0 = f32[7,3]{1,0} exponential(c0)
+  p1 = f32[3,16]{1,0} parameter(1)
+  ROOT d1 = f32[7,16]{1,0} dot(e0, p1),
+    lhs_contracting_dims={1}, rhs_contracting_dims={0}
+}
 
-  EXPECT_THAT(
-      module->entry_computation()->root_instruction(),
-      GmockMatch(
-          m::Fusion(m::Fusion(m::Parameter(), m::Parameter())
-                        .WithFusionKind(HloInstruction::FusionKind::kCustom),
-                    m::Fusion(m::Parameter(), m::Parameter())
-                        .WithFusionKind(HloInstruction::FusionKind::kCustom))
-              .WithFusionKind(HloInstruction::FusionKind::kLoop)));
+ENTRY e {
+  p0 = s8[7,3]{1,0} parameter(0)
+  p1 = f32[3,16]{1,0} parameter(1)
+  p2 = s8[7,3]{1,0} parameter(2)
+  p3 = f32[3,16]{1,0} parameter(3)
+  f0 = f32[7,16]{1,0} fusion(p0, p1), kind=kCustom, calls=triton_dot_0,
+    backend_config={"fusion_backend_config": {kind: "__triton_gemm",
+      triton_gemm_config: {"block_m":16,"block_n":16,"block_k":16,"num_stages":1,"num_warps":2,"num_ctas":1}}}
+  f1 = f32[7,16]{1,0} fusion(p2, p3), kind=kCustom, calls=triton_dot_1,
+    backend_config={"fusion_backend_config": {kind: "__triton_gemm",
+      triton_gemm_config: {"block_m":16,"block_n":16,"block_k":16,"num_stages":1,"num_warps":2,"num_ctas":1}}}
+  ROOT a = f32[7,16]{1,0} add(f0, f1)
+})";
 
-  EXPECT_TRUE(RunAndCompare(kHloText, ErrorSpec{/*aabs=*/1e-2, /*arel=*/1e-2}));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
+                       ParseAndReturnVerifiedModule(kHloText));
+  ASSERT_OK_AND_ASSIGN(bool converted, ConvertTritonGemmConfig(
+                                           device_description(), &mlir_context_)
+                                           .Run(module.get()));
+  EXPECT_TRUE(converted);
+  EXPECT_TRUE(RunAndCompareNoHloPasses(
+      std::move(module), ErrorSpec{/*aabs=*/1e-2, /*arel=*/1e-2}));
 }
 
 // TODO(b/393299275): this should just be a fusion test and does not need to be
 // in the codegen directory.
-TEST_F(TritonGemmTest, BroadcastOfScalarParameterIsFused) {
+TEST_P(TritonGemmTest, BroadcastOfScalarParameterIsFused) {
   constexpr absl::string_view kHloText = R"(
 ENTRY e {
   p0 = f16[64,256] parameter(0)
@@ -1439,7 +1445,7 @@ ENTRY e {
 
 // TODO(b/393299275): this should just be a fusion test and does not need to be
 // in the codegen directory.
-TEST_F(TritonGemmTest, BroadcastOfScalarConstantIsFused) {
+TEST_P(TritonGemmTest, BroadcastOfScalarConstantIsFused) {
   constexpr absl::string_view kHloText = R"(
 HloModule m
 
@@ -1462,7 +1468,7 @@ ENTRY e {
 
 // TODO(b/393299275): this should just be a fusion test and does not need to be
 // in the codegen directory.
-TEST_F(TritonGemmTest, DoubleBroadcastOfScalarConstantIsFused) {
+TEST_P(TritonGemmTest, DoubleBroadcastOfScalarConstantIsFused) {
   constexpr absl::string_view kHloText = R"(
 ENTRY e {
   c = s32[] constant(1)
@@ -1486,7 +1492,7 @@ ENTRY e {
 
 // TODO(b/393299275): this should just be a fusion test and does not need to be
 // in the codegen directory.
-TEST_F(TritonGemmTest, BroadcastOfVectorConstantIsFused) {
+TEST_P(TritonGemmTest, BroadcastOfVectorConstantIsFused) {
   constexpr absl::string_view kHloText = R"(
 HloModule m
 
@@ -1508,7 +1514,7 @@ ENTRY e {
 
 // TODO(b/393299275): this should just be a fusion test and does not need to be
 // in the codegen directory.
-TEST_F(TritonGemmTest, AlwaysFuseScalarConstantAtBroadcastInput) {
+TEST_P(TritonGemmTest, AlwaysFuseScalarConstantAtBroadcastInput) {
   constexpr absl::string_view kHloText = R"(
 ENTRY e {
   p0 = bf16[2,3,3]{2,1,0} parameter(0)
@@ -1536,7 +1542,7 @@ ENTRY e {
 
 // TODO(b/393299275): this should just be a fusion test and does not need to be
 // in the codegen directory.
-TEST_F(TritonGemmTest, FuseConcatenation) {
+TEST_P(TritonGemmTest, FuseConcatenation) {
   if (!SupportsBF16(GpuComputeCapability())) {
     GTEST_SKIP() << "BF16 not supported.";
   }
@@ -1566,7 +1572,7 @@ e {
 
 // TODO(b/393299275): this should just be a fusion test and does not need to be
 // in the codegen directory.
-TEST_F(TritonGemmTest, SineOutputIsFused) {
+TEST_P(TritonGemmTest, SineOutputIsFused) {
   constexpr absl::string_view kHloText = R"(
 HloModule m
 
@@ -1589,7 +1595,7 @@ ENTRY e {
 
 // TODO(b/393299275): this should just be a fusion test and does not need to be
 // in the codegen directory.
-TEST_F(TritonGemmTest, SliceInputIsFused) {
+TEST_P(TritonGemmTest, SliceInputIsFused) {
   constexpr absl::string_view kHloText = R"(
 ENTRY e {
   p0 = f16[97,121] parameter(0)
@@ -1610,7 +1616,7 @@ ENTRY e {
 
 // TODO(b/393299275): this should just be a fusion test and does not need to be
 // in the codegen directory.
-TEST_F(TritonGemmTest, SliceInputWithReshapeIsFused) {
+TEST_P(TritonGemmTest, SliceInputWithReshapeIsFused) {
   constexpr absl::string_view kHloText = R"(
 ENTRY e {
   p0 = f32[363,1536] parameter(0)
@@ -1635,7 +1641,7 @@ ENTRY e {
 // TODO(b/393299275): this should just be a fusion test and does not need to be
 // in the codegen directory. From the point of view of codegen, this is purely
 // about compositionality.
-TEST_F(TritonGemmTest, NestedSlicingWorks) {
+TEST_P(TritonGemmTest, NestedSlicingWorks) {
   constexpr absl::string_view kHloText = R"(
 ENTRY e {
   p1 = f32[6,24] parameter(1)
@@ -1657,7 +1663,7 @@ ENTRY e {
 
 // TODO(b/393299275): this should just be a fusion test and does not need to be
 // in the codegen directory.
-TEST_F(TritonGemmTest, SlicedBatchDimensionIsSupported) {
+TEST_P(TritonGemmTest, SlicedBatchDimensionIsSupported) {
   constexpr absl::string_view kHloText = R"(
 ENTRY e {
   p0 = f16[3,3,256] parameter(0)
@@ -1681,7 +1687,7 @@ ENTRY e {
   EXPECT_TRUE(RunAndCompare(kHloText, ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3}));
 }
 
-TEST_F(TritonGemmTest, BatchMajorSlicedBatchDimensionProducesCorrectResults) {
+TEST_P(TritonGemmTest, BatchMajorSlicedBatchDimensionProducesCorrectResults) {
   constexpr absl::string_view kHloText = R"(
 ENTRY e {
   p0 = f16[4,32,256] parameter(0)
@@ -1694,7 +1700,7 @@ ENTRY e {
   EXPECT_TRUE(RunAndCompare(kHloText, ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3}));
 }
 
-TEST_F(TritonGemmTest, DegenerateBatchDimensionProducesCorrectResults) {
+TEST_P(TritonGemmTest, DegenerateBatchDimensionProducesCorrectResults) {
   constexpr absl::string_view kHloText = R"(
 ENTRY e {
   p0 = f16[1,32,256] parameter(0)
@@ -1709,7 +1715,7 @@ ENTRY e {
 
 // TODO(b/393299275): this should just be a fusion test and does not need to be
 // in the codegen directory.
-TEST_F(TritonGemmTest, NarrowingConvertOutputIsFused) {
+TEST_P(TritonGemmTest, NarrowingConvertOutputIsFused) {
   constexpr absl::string_view kHloText = R"(
 HloModule m
 
@@ -1732,7 +1738,7 @@ ENTRY e {
 
 // TODO(b/393299275): this should just be a fusion test and does not need to be
 // in the codegen directory.
-TEST_F(TritonGemmTest, ParameterAfterDotIsFused) {
+TEST_P(TritonGemmTest, ParameterAfterDotIsFused) {
   if (!SupportsBF16(GpuComputeCapability())) {
     GTEST_SKIP() << "BF16 not supported.";
   }
@@ -1764,7 +1770,7 @@ ENTRY e {
   EXPECT_TRUE(RunAndCompare(kHloText, ErrorSpec{/*aabs=*/2e-2, /*arel=*/2e-2}));
 }
 
-TEST_F(TritonGemmTest, OutputFusionExecutesCorrectly) {
+TEST_P(TritonGemmTest, OutputFusionExecutesCorrectly) {
   // TODO(b/393299275): it this test useful?
   if (!SupportsBF16(GpuComputeCapability())) {
     GTEST_SKIP() << "BF16 not supported.";
@@ -1803,7 +1809,7 @@ ENTRY e {
 
 // TODO(b/393299275): this should just be a fusion test and does not need to be
 // in the codegen directory.
-TEST_F(TritonGemmTest, SplitLHSOutputTransposeAloneIsNotFused) {
+TEST_P(TritonGemmTest, SplitLHSOutputTransposeAloneIsNotFused) {
   if (!SupportsBF16(GpuComputeCapability())) {
     GTEST_SKIP() << "BF16 not supported.";
   }
@@ -1843,7 +1849,7 @@ ENTRY e {
 // time, the `bitcast` we need to hoist needs to be hoisted *upwards* but is
 // located after the dot (it collapses two consecutive non-contracting
 // dimensions together).
-TEST_F(TritonGemmTest, DISABLED_SplitLHSInputOutputIsFused) {
+TEST_P(TritonGemmTest, DISABLED_SplitLHSInputOutputIsFused) {
   if (!SupportsBF16(GpuComputeCapability())) {
     GTEST_SKIP() << "BF16 not supported.";
   }
@@ -1874,7 +1880,7 @@ ENTRY e {
   EXPECT_TRUE(RunAndCompare(kHloText, ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3}));
 }
 
-TEST_F(TritonGemmTest, SupportPredParametersUsedInExpressions) {
+TEST_P(TritonGemmTest, SupportPredParametersUsedInExpressions) {
   constexpr absl::string_view kHloText = R"(
 ENTRY e {
   p = pred[2,2]{1,0} parameter(0)
@@ -1904,7 +1910,7 @@ ENTRY e {
 // Triton emitter. This should probably be a fusion test.
 // TODO(b/353484968): Tests that don't run RunAndCompareNoHloPasses should be
 // moved to deviceless test file.
-TEST_F(TritonGemmTest, LowerDotWithLhsWithoutNonContractingDimThroughTriton) {
+TEST_P(TritonGemmTest, LowerDotWithLhsWithoutNonContractingDimThroughTriton) {
   constexpr absl::string_view kHloText = R"(
 ENTRY e {
   parameter_0 = f32[1,40] parameter(0)
@@ -1927,7 +1933,7 @@ ENTRY e {
 // Triton emitter. This should probably be a fusion test.
 // TODO(b/353484968): Tests that don't run RunAndCompareNoHloPasses should be
 // moved to deviceless test file.
-TEST_F(TritonGemmTest, LowerDotWithRhsWithoutNonContractingDimThroughTriton) {
+TEST_P(TritonGemmTest, LowerDotWithRhsWithoutNonContractingDimThroughTriton) {
   constexpr absl::string_view kHloText = R"(
 ENTRY e {
   parameter_0 = f32[1,40,250000] parameter(0)
@@ -1948,8 +1954,10 @@ ENTRY e {
 // This group of tests compares GPU results of dots already rewritten
 // into Triton fusions.
 using CompareTest = TritonGemmTest;
+INSTANTIATE_TEST_SUITE_P(All, CompareTest, ::testing::Bool(),
+                         GemmFusionVersionName);
 
-TEST_F(CompareTest, F32WithTrivialNonContractingDimension) {
+TEST_P(CompareTest, F32WithTrivialNonContractingDimension) {
   constexpr absl::string_view kHloTextRef = R"(
 HloModule r
 
@@ -2001,7 +2009,7 @@ ENTRY e {
 // runs across every possible chip that is supported by Triton. The test should
 // probably be made deviceless and repurposed to test that opt-in shared memory
 // is used only.
-TEST_F(CompareTest, UsingOptinSharedMemoryProducesSameResult) {
+TEST_P(CompareTest, UsingOptinSharedMemoryProducesSameResult) {
   if (GpuComputeCapability().IsRocm()) {
     GTEST_SKIP() << "No Optin Shared Memory on AMD.";
   }
@@ -2099,7 +2107,7 @@ ENTRY e {
                               /*run_hlo_passes=*/false));
 }
 
-TEST_F(CompareTest, NonMajorMostOutputBatchWorksCorrectly) {
+TEST_P(CompareTest, NonMajorMostOutputBatchWorksCorrectly) {
   const std::string kHloTextTest = R"(
 HloModule m
 
@@ -2167,7 +2175,7 @@ ENTRY e {
                               /*run_hlo_passes=*/false));
 }
 
-TEST_F(CompareTest, TritonDotFusionCanHaveOnlyRHSParameter) {
+TEST_P(CompareTest, TritonDotFusionCanHaveOnlyRHSParameter) {
   const std::string kHloTextTest = R"(
 HloModule m, is_scheduled=true
 
@@ -2215,7 +2223,7 @@ ENTRY e {
                                       /*run_hlo_passes=*/false));
 }
 
-TEST_F(CompareTest, TritonDotFusionCanHaveNoParametersAtAll) {
+TEST_P(CompareTest, TritonDotFusionCanHaveNoParametersAtAll) {
   const std::string kHloTextTest = R"(
 HloModule m, is_scheduled=true
 
@@ -2264,7 +2272,7 @@ ENTRY triton_gemm___computation {
                                       /*run_hlo_passes=*/false));
 }
 
-TEST_F(CompareTest, TritonDotFusionCanHaveManyParameters) {
+TEST_P(CompareTest, TritonDotFusionCanHaveManyParameters) {
   const std::string kHloTextTest = R"(
 HloModule m
 
@@ -2388,7 +2396,7 @@ ENTRY e {
                                       /*run_hlo_passes=*/false));
 }
 
-TEST_F(CompareTest, PredToBF16ConversionWorks) {
+TEST_P(CompareTest, PredToBF16ConversionWorks) {
   if (!SupportsBF16(GpuComputeCapability())) {
     GTEST_SKIP() << "BF16 not supported.";
   }
@@ -2459,7 +2467,7 @@ ENTRY e {
 // the two transposes in the fusion allows hoisting the final bitcast (`b1`).
 // I'm not sure if this is even required, since now we canonicalize transposes
 // before fusing.
-TEST_F(CompareTest, DISABLED_DifferentLayoutsAreSupportedInOneScope) {
+TEST_P(CompareTest, DISABLED_DifferentLayoutsAreSupportedInOneScope) {
   const std::string kHloTextTest = R"(
 triton_dot {
   p1 = f16[3,3,2,16]{1,3,2,0} parameter(1)
@@ -2524,7 +2532,7 @@ ENTRY e {
                                       /*run_hlo_passes=*/false));
 }
 
-TEST_F(TritonTest, UseTF32For8BitOrLessWithF32) {
+TEST_P(TritonTest, UseTF32For8BitOrLessWithF32) {
   constexpr absl::string_view kHloText = R"(
 HloModule t
 
@@ -2565,7 +2573,7 @@ CHECK:      inputPrecision = tf32
                                ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3}));
 }
 
-TEST_F(TritonGemmTest, S8ToF16DotWithSmallTileDoesNotCrash) {
+TEST_P(TritonGemmTest, S8ToF16DotWithSmallTileDoesNotCrash) {
   constexpr absl::string_view kHloText = R"(
 HloModule m
 
@@ -2593,7 +2601,7 @@ ENTRY e {
                   /*run_hlo_passes=*/false));
 }
 
-TEST_F(TritonGemmTest, S8ToF32DotWithManyWarpsDoesNotCrash) {
+TEST_P(TritonGemmTest, S8ToF32DotWithManyWarpsDoesNotCrash) {
   constexpr absl::string_view kHloText = R"(
 HloModule m
 
@@ -2621,7 +2629,7 @@ ENTRY e {
                   /*run_hlo_passes=*/false));
 }
 
-TEST_F(TritonGemmTest, Fp8DotWithSmallTileDoesNotCrash) {
+TEST_P(TritonGemmTest, Fp8DotWithSmallTileDoesNotCrash) {
   if (!GetCudaComputeCapability().IsAtLeastHopper()) {
     GTEST_SKIP() << "Doesn't pass on pre-Hopper GPUs.";
   }
@@ -2652,7 +2660,7 @@ ENTRY e {
                   /*run_hlo_passes=*/false));
 }
 
-TEST_F(TritonGemmTest, MixedF8DotExecutesCorrectly) {
+TEST_P(TritonGemmTest, MixedF8DotExecutesCorrectly) {
   if (!GetCudaComputeCapability().IsAtLeastHopper()) {
     GTEST_SKIP() << "Requires a Hopper+ GPU";
   }
@@ -2693,7 +2701,7 @@ e {
                                       /*run_hlo_passes=*/false));
 }
 
-TEST_F(TritonGemmTest, Fp8DotWithManyWarpsDoesNotCrash) {
+TEST_P(TritonGemmTest, Fp8DotWithManyWarpsDoesNotCrash) {
   if (!GetCudaComputeCapability().IsAtLeastHopper()) {
     GTEST_SKIP() << "Doesn't pass on pre-Hopper GPUs.";
   }
@@ -2726,7 +2734,7 @@ ENTRY e {
 
 // TODO(b/353484968): Tests that don't run RunAndCompareNoHloPasses should be
 // moved to deviceless test file.
-TEST_F(TritonGemmTest, WgmmaIsUsedForMemBoundShape) {
+TEST_P(TritonGemmTest, WgmmaIsUsedForMemBoundShape) {
   if (GetCudaComputeCapability().major != se::CudaComputeCapability::kHopper) {
     GTEST_SKIP() << "wgmma instruction is only available on Hopper";
   }
@@ -2764,7 +2772,7 @@ CHECK: wgmma.mma_async.sync.aligned.m64n16k16.f32.bf16.bf16
 //
 // TODO(b/353484968): Tests that don't run RunAndCompareNoHloPasses should be
 // moved to deviceless test file.
-TEST_F(TritonGemmTest, TestNoAutotuner) {
+TEST_P(TritonGemmTest, TestNoAutotuner) {
   constexpr absl::string_view kHloText = R"(
 ENTRY e {
   p0 = f16[30,30] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/collectives/BUILD b/third_party/xla/xla/backends/gpu/collectives/BUILD
index 09c9c6c0e3f109..6e6aadc1c6ff5e 100644
--- a/third_party/xla/xla/backends/gpu/collectives/BUILD
+++ b/third_party/xla/xla/backends/gpu/collectives/BUILD
@@ -1103,7 +1103,7 @@ cc_library(
         "//xla:shape_util",
         "//xla:util",
         "//xla:xla_data_proto_cc",
-        "//xla/backends/gpu/runtime:collective_kernel_api",
+        "//xla/backends/gpu/runtime:multi_gpu_barrier",
         "//xla/core/collectives:communicator",
         "//xla/core/collectives:rank_id",
         "//xla/core/collectives:reduction_kind",
diff --git a/third_party/xla/xla/backends/gpu/collectives/nccl_communicator.cc b/third_party/xla/xla/backends/gpu/collectives/nccl_communicator.cc
index 43d9f3b7073205..cad31b7fe963e4 100644
--- a/third_party/xla/xla/backends/gpu/collectives/nccl_communicator.cc
+++ b/third_party/xla/xla/backends/gpu/collectives/nccl_communicator.cc
@@ -49,7 +49,7 @@ limitations under the License.
 #include "xla/backends/gpu/collectives/nccl_symmetric_memory.h"
 #include "xla/backends/gpu/collectives/nccl_types.h"
 #include "xla/backends/gpu/collectives/single_threaded_executor.h"
-#include "xla/backends/gpu/runtime/collective_kernel_api.h"
+#include "xla/backends/gpu/runtime/multi_gpu_barrier.h"
 #include "xla/core/collectives/communicator.h"
 #include "xla/core/collectives/rank_id.h"
 #include "xla/core/collectives/reduction_kind.h"
diff --git a/third_party/xla/xla/backends/gpu/runtime/BUILD b/third_party/xla/xla/backends/gpu/runtime/BUILD
index 75bdcc99a3ef38..53b24c32bff7be 100644
--- a/third_party/xla/xla/backends/gpu/runtime/BUILD
+++ b/third_party/xla/xla/backends/gpu/runtime/BUILD
@@ -2360,12 +2360,12 @@ cc_library(
     deps = [
         ":collective_clique_requests",
         ":collective_cliques",
-        ":collective_kernel_api",
         ":collective_memory",
         ":collective_memory_requests",
         ":collective_thunk",
         ":collective_thunk_proto_cc",
         ":command_state",
+        ":multi_gpu_barrier",
         ":ragged_all_to_all",
         ":thunk",
         ":thunk_proto_cc",
@@ -2946,8 +2946,8 @@ cc_library(
     compatible_with = get_compatible_with_portable(),
     deps = [
         ":collective_clique_requests",
-        ":collective_kernel_api",
         ":collective_params",
+        ":multi_gpu_barrier",
         "//xla:status_macros",
         "//xla:util",
         "//xla/backends/gpu/collectives:gpu_clique",
@@ -3431,9 +3431,9 @@ cc_library(
 )
 
 cc_library(
-    name = "collective_kernel_api",
-    srcs = ["collective_kernel_api.cc"],
-    hdrs = ["collective_kernel_api.h"],
+    name = "multi_gpu_barrier",
+    srcs = ["multi_gpu_barrier.cc"],
+    hdrs = ["multi_gpu_barrier.h"],
     compatible_with = get_compatible_with_portable(),
     deps = [
         "//xla:status_macros",
@@ -3449,7 +3449,6 @@ cc_library(
         "@com_google_absl//absl/base:no_destructor",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/status",
-        "@com_google_absl//absl/status:status_macros",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/synchronization",
     ],
diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_cliques.cc b/third_party/xla/xla/backends/gpu/runtime/collective_cliques.cc
index 3855c93ee98768..b0fb8788dc98d8 100644
--- a/third_party/xla/xla/backends/gpu/runtime/collective_cliques.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/collective_cliques.cc
@@ -36,8 +36,8 @@ limitations under the License.
 #include "xla/backends/gpu/collectives/gpu_cliques.h"
 #include "xla/backends/gpu/collectives/gpu_communicator.h"
 #include "xla/backends/gpu/runtime/collective_clique_requests.h"
-#include "xla/backends/gpu/runtime/collective_kernel_api.h"
 #include "xla/backends/gpu/runtime/collective_params.h"
+#include "xla/backends/gpu/runtime/multi_gpu_barrier.h"
 #include "xla/core/collectives/clique_id.h"
 #include "xla/core/collectives/clique_key.h"
 #include "xla/core/collectives/communicator.h"
diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_api.cc b/third_party/xla/xla/backends/gpu/runtime/multi_gpu_barrier.cc
similarity index 98%
rename from third_party/xla/xla/backends/gpu/runtime/collective_kernel_api.cc
rename to third_party/xla/xla/backends/gpu/runtime/multi_gpu_barrier.cc
index 0eb6811f600d7b..4f2ce8c8f39aa1 100644
--- a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_api.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/multi_gpu_barrier.cc
@@ -13,7 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include "xla/backends/gpu/runtime/collective_kernel_api.h"
+#include "xla/backends/gpu/runtime/multi_gpu_barrier.h"
 
 #include <algorithm>
 #include <array>
diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_api.h b/third_party/xla/xla/backends/gpu/runtime/multi_gpu_barrier.h
similarity index 81%
rename from third_party/xla/xla/backends/gpu/runtime/collective_kernel_api.h
rename to third_party/xla/xla/backends/gpu/runtime/multi_gpu_barrier.h
index cd601d10610dd9..c38710e17e15ff 100644
--- a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_api.h
+++ b/third_party/xla/xla/backends/gpu/runtime/multi_gpu_barrier.h
@@ -13,8 +13,8 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#ifndef XLA_BACKENDS_GPU_RUNTIME_COLLECTIVE_KERNEL_API_H_
-#define XLA_BACKENDS_GPU_RUNTIME_COLLECTIVE_KERNEL_API_H_
+#ifndef XLA_BACKENDS_GPU_RUNTIME_MULTI_GPU_BARRIER_H_
+#define XLA_BACKENDS_GPU_RUNTIME_MULTI_GPU_BARRIER_H_
 
 #include <cstddef>
 #include <cstdint>
@@ -28,12 +28,6 @@ limitations under the License.
 
 namespace xla::gpu {
 
-// This file contains collective kernels related method shared between XLA and
-// Mosaic. Be careful when changing the API.
-//
-// Collective kernels are kernels which are launched simultaneously on multiple
-// devices, and have an access to each other device's memory.
-
 // Launches a cross-GPU barrier synchronization.
 absl::Status LaunchMultiGpuBarrier(
     stream_executor::Stream* stream, int64_t num_devices, RankId rank,
@@ -57,4 +51,4 @@ size_t GetMultiGpuBarrierSignalValueSize();
 
 }  // namespace xla::gpu
 
-#endif  // XLA_BACKENDS_GPU_RUNTIME_COLLECTIVE_KERNEL_API_H_
+#endif  // XLA_BACKENDS_GPU_RUNTIME_MULTI_GPU_BARRIER_H_
diff --git a/third_party/xla/xla/backends/gpu/runtime/ragged_all_to_all_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/ragged_all_to_all_thunk.cc
index 547dea71fffb9b..4c4b60f7cf0277 100644
--- a/third_party/xla/xla/backends/gpu/runtime/ragged_all_to_all_thunk.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/ragged_all_to_all_thunk.cc
@@ -48,12 +48,12 @@ limitations under the License.
 #include "xla/backends/gpu/collectives/gxl_communicator.h"
 #include "xla/backends/gpu/runtime/collective_clique_requests.h"
 #include "xla/backends/gpu/runtime/collective_cliques.h"
-#include "xla/backends/gpu/runtime/collective_kernel_api.h"
 #include "xla/backends/gpu/runtime/collective_memory.h"
 #include "xla/backends/gpu/runtime/collective_memory_requests.h"
 #include "xla/backends/gpu/runtime/collective_thunk.h"
 #include "xla/backends/gpu/runtime/collective_thunk.pb.h"
 #include "xla/backends/gpu/runtime/command_state.h"
+#include "xla/backends/gpu/runtime/multi_gpu_barrier.h"
 #include "xla/backends/gpu/runtime/ragged_all_to_all.h"
 #include "xla/backends/gpu/runtime/thunk.h"
 #include "xla/backends/gpu/runtime/thunk.pb.h"
diff --git a/third_party/xla/xla/backends/gpu/transforms/conv_fusion_rewriter.cc b/third_party/xla/xla/backends/gpu/transforms/conv_fusion_rewriter.cc
index 7cef77dadcac8a..7022d28feceac9 100644
--- a/third_party/xla/xla/backends/gpu/transforms/conv_fusion_rewriter.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/conv_fusion_rewriter.cc
@@ -82,6 +82,15 @@ bool IsConvFusionOutputsValid(const std::vector<HloInstruction*>& outputs) {
                                  (outputs[1]->opcode() == HloOpcode::kReduce)) {
     return false;
   }
+  // Disallow upcast converts at fusion outputs. S32->F32 is allowed as cuDNN
+  // INT8 convs require it for epilogue fusions.
+  for (const HloInstruction* output : outputs) {
+    if (output->opcode() == HloOpcode::kConvert &&
+        output->shape().element_type() == F32 &&
+        output->operand(0)->shape().element_type() != S32) {
+      return false;
+    }
+  }
   return true;
 }
 
diff --git a/third_party/xla/xla/backends/gpu/transforms/conv_fusion_rewriter_test.cc b/third_party/xla/xla/backends/gpu/transforms/conv_fusion_rewriter_test.cc
index 7cf3fa00701406..ac6e739d45f2d4 100644
--- a/third_party/xla/xla/backends/gpu/transforms/conv_fusion_rewriter_test.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/conv_fusion_rewriter_test.cc
@@ -575,6 +575,38 @@ TEST_F(ConvFusionRewriterUnitTest, Test1DBiasBroadcastFusedF16) {
                   .WithShape(F16, {1, 9, 9, 32}));
 }
 
+TEST_F(ConvFusionRewriterUnitTest, Test1DBiasBroadcastSharedWithMultipleConvs) {
+  RunAndMatch(
+      R"(
+    HloModule Test
+
+    ENTRY Test {
+      input1 = f16[1,9,9,17] parameter(0)
+      filter1 = f16[32,3,3,17] parameter(1)
+      filter2 = f16[32,3,3,32] parameter(2)
+      bias = f16[32] parameter(3)
+      bias_broadcast = f16[1,9,9,32] broadcast(bias), dimensions={3}
+      zero = f16[] constant(0)
+      zeros = f16[1,9,9,32] broadcast(zero), dimensions={}
+
+      conv1 = f16[1,9,9,32] convolution(input1, filter1),
+                window={size=3x3 pad=1_1x1_1},
+                dim_labels=b01f_o01i->b01f
+      sum1 = add(conv1, bias_broadcast)
+      relu1 = maximum(sum1, zeros)
+
+      conv2 = f16[1,9,9,32] convolution(relu1, filter2),
+                window={size=3x3 pad=1_1x1_1},
+                dim_labels=b01f_o01i->b01f
+      sum2 = add(conv2, bias_broadcast)
+      ROOT relu2 = maximum(sum2, zeros)
+    })",
+      m::Fusion(m::Fusion(m::Parameter(0), m::Parameter(1), m::Parameter(3)),
+                m::Parameter(2), m::Parameter(3))
+          .WithFusionKind(HloInstruction::FusionKind::kCustom)
+          .WithShape(F16, {1, 9, 9, 32}));
+}
+
 TEST_F(ConvFusionRewriterUnitTest, FuseAlpha) {
   MAYBE_SKIP_TEST("I8");
   RunAndMatch(R"(
diff --git a/third_party/xla/xla/backends/gpu/transforms/cudnn_fusion_compiler.cc b/third_party/xla/xla/backends/gpu/transforms/cudnn_fusion_compiler.cc
index d20b178c722844..3475a14ab92ef8 100644
--- a/third_party/xla/xla/backends/gpu/transforms/cudnn_fusion_compiler.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/cudnn_fusion_compiler.cc
@@ -529,6 +529,31 @@ class ConvDimensionAdapter {
     return -1;
   }
 
+  const HloInstruction* get_broadcast_user(const HloInstruction* hlo) {
+    auto all_users_are_broadcast = [](const HloInstruction* instr) {
+      return !instr->users().empty() &&
+             absl::c_all_of(instr->users(), [](const HloInstruction* u) {
+               return u->opcode() == HloOpcode::kBroadcast;
+             });
+    };
+
+    // Pattern 1: hlo -> broadcast
+    if (all_users_are_broadcast(hlo)) {
+      return hlo->users()[0];
+    }
+
+    // Pattern 2: hlo -> convert -> broadcast
+    if (hlo->user_count() == 1 &&
+        hlo->users()[0]->opcode() == HloOpcode::kConvert) {
+      const HloInstruction* convert = hlo->users()[0];
+      if (all_users_are_broadcast(convert)) {
+        return convert->users()[0];
+      }
+    }
+
+    return nullptr;
+  };
+
   std::optional<Result> DimensionsAndStrides(const HloInstruction& hlo) {
     int64_t spatial_dims =
         std::max<int64_t>(2, dums_.input_spatial_dimensions_size());
@@ -549,9 +574,8 @@ class ConvDimensionAdapter {
 
       // If the parameter is consumed by a broadcast, map its dimensions to the
       // corresponding cuDNN canonical axes (N, C, spatial...).
-      if (hlo.user_count() == 1 &&
-          hlo.users()[0]->opcode() == HloOpcode::kBroadcast) {
-        const auto& bcast_dims = hlo.users()[0]->dimensions();
+      if (const HloInstruction* broadcast = get_broadcast_user(&hlo)) {
+        const auto& bcast_dims = broadcast->dimensions();
         for (int i = 0; i < bcast_dims.size(); ++i) {
           int64_t cudnn_dim = HloDimToCudnnDim(bcast_dims[i]);
           if (cudnn_dim >= 0 && cudnn_dim < cudnn_rank) {
@@ -559,15 +583,18 @@ class ConvDimensionAdapter {
             result.strides[cudnn_dim] = 1;
           }
         }
-      } else if (hlo.shape().dimensions().size() == 1) {
-        // Fallback for un-broadcasted 1D parameters: assume channel bias [1, C,
-        // 1, 1].
+        return result;
+      }
+
+      // Fallback for un-broadcasted 1D parameters: assume channel bias [1, C,
+      // 1, 1].
+      if (hlo.shape().dimensions().size() == 1) {
         result.sizes[1] = hlo.shape().dimensions(0);
         result.strides[1] = 1;
-      } else {
-        return std::nullopt;
+        return result;
       }
-      return result;
+
+      return std::nullopt;
     }
     // Placeholder FP32 data type here, it is not used.
     auto desc = se::dnn::TensorDescriptor::For(
@@ -831,10 +858,14 @@ absl::StatusOr<se::gpu::CudnnGraph> HloFusionToCuDnnGraph(
       // and int32 = conv(int8, int8)
       hlo_to_cudnn[hlo] = operand(0);
       continue;
-    } else if (HloPredicateIsOp<HloOpcode::kParameter>(hlo)) {
+    }
+
+    if (HloPredicateIsOp<HloOpcode::kParameter>(hlo)) {
       CHECK(hlo_to_cudnn.contains(hlo));
       continue;
-    } else if (HloPredicateIsOp<HloOpcode::kCustomCall>(hlo)) {
+    }
+
+    if (HloPredicateIsOp<HloOpcode::kCustomCall>(hlo)) {
       if (hlo->user_count() != 1 ||
           !IsWorkspaceAllocationRoot(*hlo->users()[0])) {
         return absl::UnimplementedError(
@@ -843,7 +874,9 @@ absl::StatusOr<se::gpu::CudnnGraph> HloFusionToCuDnnGraph(
                          hlo->ToString()));
       }
       continue;
-    } else if (HloPredicateIsOp<HloOpcode::kTuple>(hlo)) {
+    }
+
+    if (HloPredicateIsOp<HloOpcode::kTuple>(hlo)) {
       if (!IsWorkspaceAllocationRoot(*hlo) && !IsAmaxRoot(*hlo)) {
         return absl::UnimplementedError(
             absl::StrCat("Tuples are only expected at outputs for workspace "
@@ -851,9 +884,17 @@ absl::StatusOr<se::gpu::CudnnGraph> HloFusionToCuDnnGraph(
                          hlo->ToString()));
       }
       continue;
-    } else if (HloPredicateIsOp<HloOpcode::kConstant>(hlo)) {
+    }
+
+    if (HloPredicateIsOp<HloOpcode::kConstant>(hlo)) {
+      const Shape& root_shape = computation.root_instruction()->shape();
+      const Shape& output_shape =
+          root_shape.IsTuple() ? root_shape.tuple_shapes(0) : root_shape;
+      int64_t rank = std::max<int64_t>(3, output_shape.dimensions().size());
+      ABSL_ASSIGN_OR_RETURN(hlo_to_cudnn[hlo],
+                       HandleConstantHloToCudnnGraph(*hlo, graph, rank));
       ABSL_ASSIGN_OR_RETURN(hlo_to_cudnn[hlo],
-                       HandleConstantHloToCudnnGraph(*hlo, graph));
+                       HandleConstantHloToCudnnGraph(*hlo, graph, rank));
     } else if (HloPredicateIsOp<HloOpcode::kReshape, HloOpcode::kBitcast,
                                 HloOpcode::kTranspose, HloOpcode::kCopy,
                                 HloOpcode::kSlice>(hlo)) {
diff --git a/third_party/xla/xla/backends/gpu/transforms/cudnn_fusion_utils.cc b/third_party/xla/xla/backends/gpu/transforms/cudnn_fusion_utils.cc
index a60ab8e6ad60fb..dd4b4d18606f3b 100644
--- a/third_party/xla/xla/backends/gpu/transforms/cudnn_fusion_utils.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/cudnn_fusion_utils.cc
@@ -151,6 +151,25 @@ bool IsEpilogueOpSupportedByCuDNN(const HloInstruction& hlo,
   }
 }
 
+namespace {
+
+bool IsCheapToDuplicate(const HloInstruction* hlo) {
+  if (hlo->opcode() == HloOpcode::kConstant) {
+    return ShapeUtil::IsScalar(hlo->shape());
+  }
+  if (hlo->opcode() == HloOpcode::kBroadcast) {
+    return ShapeUtil::IsScalar(hlo->operand(0)->shape()) ||
+           hlo->operand(0)->shape().dimensions().size() == 1;
+  }
+  if (hlo->opcode() == HloOpcode::kConvert) {
+    return ShapeUtil::IsScalar(hlo->shape()) ||
+           hlo->shape().dimensions().size() == 1;
+  }
+  return false;
+}
+
+}  // namespace
+
 HloInstruction* FuseTowardOperand(
     HloInstruction* hlo, HloComputation::Builder& builder,
     std::vector<HloInstruction*>& fusion_params,
@@ -162,7 +181,7 @@ HloInstruction* FuseTowardOperand(
   HloInstruction* fused_hlo;
   if (IsEpilogueOpSupportedByCuDNN(*hlo, /*can_fuse_reduce=*/false, is_nchw,
                                    device_info) &&
-      hlo->user_count() == 1) {
+      (hlo->user_count() == 1 || IsCheapToDuplicate(hlo))) {
     HloInstruction::InstructionVector new_operands;
     for (int i = 0; i < hlo->operand_count(); ++i) {
       HloInstruction* operand = hlo->mutable_operand(i);
diff --git a/third_party/xla/xla/backends/gpu/transforms/gemm_fusion_test.cc b/third_party/xla/xla/backends/gpu/transforms/gemm_fusion_test.cc
index c1c6e89ad2b2bc..7022507b428b44 100644
--- a/third_party/xla/xla/backends/gpu/transforms/gemm_fusion_test.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/gemm_fusion_test.cc
@@ -78,6 +78,7 @@ class GemmFusionTestBase : public HloHardwareIndependentTestBase {
         HloHardwareIndependentTestBase::GetDebugOptionsForTest();
     debug_options.set_xla_gpu_triton_gemm_any(false);
     debug_options.set_xla_gpu_gemm_rewrite_size_threshold(0);
+    debug_options.set_xla_gpu_experimental_gemm_fusion_v2(false);
     return debug_options;
   }
 
@@ -96,15 +97,16 @@ class GemmFusionTest : public GemmFusionTestBase,
  public:
   DebugOptions GetDebugOptionsForTest() const override {
     DebugOptions debug_options = GemmFusionTestBase::GetDebugOptionsForTest();
-    if (GetParam()) {
-      debug_options.set_xla_gpu_experimental_enable_tiling_propagation(true);
-    }
+    debug_options.set_xla_gpu_experimental_enable_tiling_propagation(
+        GetParam());
     return debug_options;
   }
 };
 
-// Create a parameterized test that makes sure that both the legacy and the new
-// implementation of dot fusion are working as expected.
+// While we launch, create a parameterized test to test all combinations. This
+// test class has 2 parameters:
+// 1. Whether to use Gemm Fusion V1 or V2.
+// 2. Whether to use symbolic analysis or tiling propagation.
 class GemmFusionTestVersioned
     : public GemmFusionTestBase,
       public ::testing::WithParamInterface<std::tuple<bool, bool>> {
@@ -113,9 +115,8 @@ class GemmFusionTestVersioned
     DebugOptions debug_options = GemmFusionTestBase::GetDebugOptionsForTest();
     debug_options.set_xla_gpu_experimental_gemm_fusion_v2(
         std::get<0>(GetParam()));
-    if (std::get<1>(GetParam())) {
-      debug_options.set_xla_gpu_experimental_enable_tiling_propagation(true);
-    }
+    debug_options.set_xla_gpu_experimental_enable_tiling_propagation(
+        std::get<1>(GetParam()));
     return debug_options;
   }
 };
@@ -128,9 +129,8 @@ class GemmFusionTestV2 : public GemmFusionTestBase,
   DebugOptions GetDebugOptionsForTest() const override {
     DebugOptions debug_options = GemmFusionTestBase::GetDebugOptionsForTest();
     debug_options.set_xla_gpu_experimental_gemm_fusion_v2(true);
-    if (GetParam()) {
-      debug_options.set_xla_gpu_experimental_enable_tiling_propagation(true);
-    }
+    debug_options.set_xla_gpu_experimental_enable_tiling_propagation(
+        GetParam());
     return debug_options;
   }
 };
diff --git a/third_party/xla/xla/backends/gpu/transforms/hoist_fused_bitcasts_test.cc b/third_party/xla/xla/backends/gpu/transforms/hoist_fused_bitcasts_test.cc
index 9bec365d533c46..ddbd4ca0ec34e4 100644
--- a/third_party/xla/xla/backends/gpu/transforms/hoist_fused_bitcasts_test.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/hoist_fused_bitcasts_test.cc
@@ -76,6 +76,14 @@ class HoistFusedBitcastsReshapeTest
   HoistFusedBitcastsReshapeTest() {
     RegisterSymbolicExprStorage(&mlir_context_);
   }
+
+  DebugOptions GetDebugOptionsForTest() const override {
+    DebugOptions debug_options =
+        HloHardwareIndependentTestBase::GetDebugOptionsForTest();
+    debug_options.set_xla_gpu_experimental_gemm_fusion_v2(false);
+    return debug_options;
+  }
+
   const se::DeviceDescription device_description_{
       TestGpuDeviceInfo::RTXA6000DeviceInfo(
           se::GpuComputeCapability{se::CudaComputeCapability::Ampere()})};
diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_computation_deduplicator.cc b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_computation_deduplicator.cc
index de2bc2387b4c91..04a6526c803318 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_computation_deduplicator.cc
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_computation_deduplicator.cc
@@ -91,9 +91,12 @@ absl::StatusOr<bool> HloComputationDeduplicator::RunImpl(
        module->MakeComputationPostOrder(execution_threads)) {
     // Ignore entry computation since it is called from outside and computations
     // with large number of instructions or large-size constants due to increase
-    // in time taken to stringify.
+    // in time taken to stringify. Also ignore fusion computations, which need
+    // to have a 1:1 relationship between caller and computation, except if we
+    // run in the mode where we want to annotate duplicate fusion computations.
     if (comp->IsEntryComputation() || comp->instruction_count() > 128 ||
-        ContainsLargeConstants(comp)) {
+        ContainsLargeConstants(comp) ||
+        (!mark_fusion_duplications_ && comp->IsFusionComputation())) {
       continue;
     }
     // Don't deduplicate collectives and non-collectives.
diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_computation_deduplicator_test.cc b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_computation_deduplicator_test.cc
index 9e0b3c860400f9..4b668a3593e09f 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_computation_deduplicator_test.cc
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_computation_deduplicator_test.cc
@@ -21,6 +21,7 @@ limitations under the License.
 #include <string>
 #include <vector>
 
+#include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include "absl/strings/string_view.h"
 #include "xla/hlo/ir/hlo_computation.h"
@@ -629,7 +630,7 @@ TEST_F(HloComputationDeduplicatorTest, LargeSubComputationTest) {
   }
   module->AddEntryComputation(main.Build());
   HloComputationDeduplicator dedup;
-  TF_ASSERT_OK_AND_ASSIGN(bool changed, dedup.Run(module.get()));
+  ASSERT_OK_AND_ASSIGN(bool changed, dedup.Run(module.get()));
   EXPECT_FALSE(changed);
   std::vector<HloComputation *> computations = module->MakeComputationSorted();
   EXPECT_EQ(computations.size(), (total_regions + 1));
@@ -665,8 +666,8 @@ TEST_F(HloComputationDeduplicatorTest, DontDeduplicateReduceAllReduce) {
 }
 
 TEST_F(HloComputationDeduplicatorTest, DeduplicateChain) {
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
-                          ParseAndReturnVerifiedModule(R"(
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(R"(
 HloModule module
 
 fusion0 {
@@ -696,8 +697,13 @@ ENTRY entry {
   ROOT add = f32[] add(fusion.0, p1)
 }
 )"));
+  HloComputationDeduplicator dedup_do_not_mark(
+      /*mark_fusion_duplications=*/false);
+  ASSERT_OK_AND_ASSIGN(bool changed, dedup_do_not_mark.Run(module.get()));
+  EXPECT_FALSE(changed);
+
   HloComputationDeduplicator dedup(/*mark_fusion_duplications=*/true);
-  TF_ASSERT_OK_AND_ASSIGN(bool changed, dedup.Run(module.get()));
+  ASSERT_OK_AND_ASSIGN(changed, dedup.Run(module.get()));
   EXPECT_TRUE(changed);
   EXPECT_EQ(module->computation_count(), 4);
   HloInstruction* fusion0 =
diff --git a/third_party/xla/xla/python/ifrt/remap_plan.cc b/third_party/xla/xla/python/ifrt/remap_plan.cc
index e0e43c814425e4..4ea7e34451f3ce 100644
--- a/third_party/xla/xla/python/ifrt/remap_plan.cc
+++ b/third_party/xla/xla/python/ifrt/remap_plan.cc
@@ -576,6 +576,14 @@ absl::Status RemapPlan::Validate() const {
     }
   }
 
+  if (!rep_->input_devices_for_output_map.empty() &&
+      rep_->input_devices_for_output_map.size() != num_outputs) {
+    return InvalidArgument(
+        "`input_devices_for_output_map` has %d outputs, but expected %d "
+        "outputs",
+        rep_->input_devices_for_output_map.size(), num_outputs);
+  }
+
   for (int i = 0; i < num_outputs; ++i) {
     xla::ifrt::DeviceList* devices =
         rep_->output_specs[i].sharding->devices()->AddressableDeviceList();
diff --git a/third_party/xla/xla/python/ifrt/remap_plan_test.cc b/third_party/xla/xla/python/ifrt/remap_plan_test.cc
index 83bfe931cb90e5..a32fe7ffd867c0 100644
--- a/third_party/xla/xla/python/ifrt/remap_plan_test.cc
+++ b/third_party/xla/xla/python/ifrt/remap_plan_test.cc
@@ -735,6 +735,29 @@ TEST_P(RemapPlanTest, InvalidInputDevicesForOutputMap) {
                                HasSubstr("does not reference that device")));
   }
 
+  {
+    std::vector<ArraySpec> multi_output_specs = {dummy_spec, dummy_spec};
+    std::vector<RemapPlan::Mapping> multi_mappings = {
+        RemapPlan::Mapping{/*in_array=*/0,
+                           /*out_array=*/0,
+                           /*from=*/{RemapPlan::Interval{0, 2, 1}},
+                           /*to=*/{RemapPlan::Interval{0, 2, 1}}},
+        RemapPlan::Mapping{/*in_array=*/1,
+                           /*out_array=*/1,
+                           /*from=*/{RemapPlan::Interval{0, 2, 1}},
+                           /*to=*/{RemapPlan::Interval{0, 2, 1}}}};
+    absl::flat_hash_map<int, std::vector<RemapPlan::InputDeviceRange>>
+        input_devices_for_output_map;
+    input_devices_for_output_map.insert(
+        {0, {{0, dummy_spec.sharding->devices()}}});
+    RemapPlan plan(input_specs, multi_output_specs, multi_mappings,
+                   std::move(input_devices_for_output_map));
+    EXPECT_THAT(plan.Validate(),
+                absl_testing::StatusIs(
+                    absl::StatusCode::kInvalidArgument,
+                    HasSubstr("has 1 outputs, but expected 2 outputs")));
+  }
+
   ASSERT_OK(RemapPlan::CreateOptimized(client(), std::move(input_specs),
                                        std::move(output_specs),
                                        std::move(mappings))
diff --git a/third_party/xla/xla/python/pjrt_ifrt/pjrt_remap.cc b/third_party/xla/xla/python/pjrt_ifrt/pjrt_remap.cc
index 9ae4d055a26570..bbddb6ed5934c8 100644
--- a/third_party/xla/xla/python/pjrt_ifrt/pjrt_remap.cc
+++ b/third_party/xla/xla/python/pjrt_ifrt/pjrt_remap.cc
@@ -17,9 +17,11 @@ limitations under the License.
 
 #include <cstdint>
 #include <memory>
+#include <optional>
 #include <utility>
 #include <vector>
 
+#include "absl/container/flat_hash_map.h"
 #include "absl/log/check.h"
 #include "absl/status/status_macros.h"
 #include "absl/status/statusor.h"
@@ -27,15 +29,15 @@ limitations under the License.
 #include "xla/pjrt/pjrt_client.h"
 #include "xla/pjrt/pjrt_layout.h"
 #include "xla/python/ifrt/array.h"
+#include "xla/python/ifrt/device.h"
+#include "xla/python/ifrt/device_list.h"
 #include "xla/python/ifrt/dtype.h"
 #include "xla/python/ifrt/remap_plan.h"
 #include "xla/python/ifrt/rtti.h"
 #include "xla/python/ifrt/shape.h"
 #include "xla/python/pjrt_ifrt/pjrt_array.h"
 #include "xla/tsl/concurrency/ref_count.h"
-#include "xla/tsl/platform/errors.h"
 #include "xla/tsl/platform/logging.h"
-#include "xla/tsl/platform/statusor.h"
 #include "xla/util.h"
 
 namespace xla {
@@ -97,34 +99,160 @@ PjRtCompatibleClientRemapArrays(PjRtCompatibleClient* client,
                                    ->size());
   }
 
-  for (const RemapPlan::Mapping& mapping : plan.mappings()) {
-    ABSL_ASSIGN_OR_RETURN(
-        absl::Span<std::shared_ptr<xla::PjRtBuffer>> in_buffers,
-        static_cast<PjRtCompatibleArray*>(arrays[mapping.in_array].get())
-            ->mutable_pjrt_buffers());
-    PjRtArray::PjRtBuffers& out_buffers = out_buffers_list[mapping.out_array];
-    for (int s = 0; s < mapping.from.size(); ++s) {
-      const RemapPlan::Interval& in_interval = mapping.from[s];
-      const RemapPlan::Interval& out_interval = mapping.to[s];
-      int64_t in_shard = in_interval.start;
-      int64_t out_shard = out_interval.start;
-      switch (semantics) {
-        case ArrayCopySemantics::kReuseInput:
-          while (in_shard < in_interval.end) {
-            out_buffers[out_shard] = in_buffers[in_shard];
-            in_shard += in_interval.step;
-            out_shard += out_interval.step;
+  // Handle outputs using `input_devices_for_output_map` when specified.
+  if (!plan.input_devices_for_output_map().empty()) {
+    std::vector<std::optional<absl::flat_hash_map<Device*, int>>>
+        in_device_to_shard(num_inputs);
+
+    for (const auto& [out_array, input_ranges] :
+         // NOLINTNEXTLINE(*-custom-deterministic-iteration-order)
+         plan.input_devices_for_output_map()) {
+      if (out_array < 0 || out_array >= num_outputs) {
+        return InvalidArgument("out_array must be in [0, %d), but is %d",
+                               num_outputs, out_array);
+      }
+      absl::Span<Device* const> out_devices = plan.output_specs()[out_array]
+                                                  .sharding->devices()
+                                                  ->AddressableDeviceList()
+                                                  ->devices();
+      // Skip outputs with no addressable shards on this controller.
+      if (out_devices.empty()) {
+        continue;
+      }
+
+      absl::flat_hash_map<Device*, int> device_to_out_shard;
+      device_to_out_shard.reserve(out_devices.size());
+      for (int out_shard = 0; out_shard < out_devices.size(); ++out_shard) {
+        device_to_out_shard.insert({out_devices[out_shard], out_shard});
+      }
+
+      PjRtArray::PjRtBuffers& out_buffers = out_buffers_list[out_array];
+
+      for (const RemapPlan::InputDeviceRange& input_range : input_ranges) {
+        absl::Span<Device* const> range_devices =
+            input_range.input_devices->AddressableDeviceList()->devices();
+        // Skip input ranges with no addressable shards on this controller.
+        if (range_devices.empty()) {
+          continue;
+        }
+
+        int in_array = input_range.in_array;
+        if (in_array < 0 || in_array >= num_inputs) {
+          return InvalidArgument("in_array must be in [0, %d), but is %d",
+                                 num_inputs, in_array);
+        }
+        ABSL_ASSIGN_OR_RETURN(
+            absl::Span<std::shared_ptr<xla::PjRtBuffer>> in_buffers,
+            static_cast<PjRtCompatibleArray*>(arrays[in_array].get())
+                ->mutable_pjrt_buffers());
+
+        if (!in_device_to_shard[in_array].has_value()) {
+          absl::Span<Device* const> in_devices = arrays[in_array]
+                                                     ->sharding()
+                                                     .devices()
+                                                     ->AddressableDeviceList()
+                                                     ->devices();
+          auto& map = in_device_to_shard[in_array].emplace();
+          map.reserve(in_devices.size());
+          for (int in_shard = 0; in_shard < in_devices.size(); ++in_shard) {
+            map.insert({in_devices[in_shard], in_shard});
+          }
+        }
+
+        for (Device* const device : range_devices) {
+          auto in_it = in_device_to_shard[in_array]->find(device);
+          if (in_it == in_device_to_shard[in_array]->end()) {
+            return InvalidArgument("Device %s not found in input array %d",
+                                   device->DebugString(), in_array);
           }
-          break;
-        case ArrayCopySemantics::kDonateInput:
-          while (in_shard < in_interval.end) {
-            out_buffers[out_shard] = std::move(in_buffers[in_shard]);
-            in_shard += in_interval.step;
-            out_shard += out_interval.step;
+          int in_shard = in_it->second;
+
+          auto out_it = device_to_out_shard.find(device);
+          if (out_it == device_to_out_shard.end()) {
+            return InvalidArgument("Device %s not found in output array %d",
+                                   device->DebugString(), out_array);
+          }
+          int out_shard = out_it->second;
+
+          if (out_buffers[out_shard] != nullptr) {
+            return InvalidArgument(
+                "Output array %d shard on device %s is assigned more than once",
+                out_array, device->DebugString());
           }
-          break;
-        default:
-          return InvalidArgument("Invalid ArrayCopySemantics: %d", semantics);
+
+          switch (semantics) {
+            case ArrayCopySemantics::kReuseInput:
+              out_buffers[out_shard] = in_buffers[in_shard];
+              break;
+            case ArrayCopySemantics::kDonateInput:
+              if (in_buffers[in_shard] == nullptr) {
+                return InvalidArgument(
+                    "Input array %d shard on device %s is used more than once",
+                    in_array, device->DebugString());
+              }
+              out_buffers[out_shard] = std::move(in_buffers[in_shard]);
+              break;
+            default:
+              return InvalidArgument("Invalid ArrayCopySemantics: %d",
+                                     semantics);
+          }
+        }
+      }
+    }
+  } else {
+    for (const RemapPlan::Mapping& mapping : plan.mappings()) {
+      ABSL_ASSIGN_OR_RETURN(
+          absl::Span<std::shared_ptr<xla::PjRtBuffer>> in_buffers,
+          static_cast<PjRtCompatibleArray*>(arrays[mapping.in_array].get())
+              ->mutable_pjrt_buffers());
+      PjRtArray::PjRtBuffers& out_buffers = out_buffers_list[mapping.out_array];
+      for (int s = 0; s < mapping.from.size(); ++s) {
+        const RemapPlan::Interval& in_interval = mapping.from[s];
+        const RemapPlan::Interval& out_interval = mapping.to[s];
+        int64_t in_shard = in_interval.start;
+        int64_t out_shard = out_interval.start;
+        switch (semantics) {
+          case ArrayCopySemantics::kReuseInput:
+            while (in_shard < in_interval.end) {
+              if (out_buffers[out_shard] != nullptr) {
+                return InvalidArgument(
+                    "Output array %d shard %d is assigned more than once",
+                    mapping.out_array, out_shard);
+              }
+              out_buffers[out_shard] = in_buffers[in_shard];
+              in_shard += in_interval.step;
+              out_shard += out_interval.step;
+            }
+            break;
+          case ArrayCopySemantics::kDonateInput:
+            while (in_shard < in_interval.end) {
+              if (out_buffers[out_shard] != nullptr) {
+                return InvalidArgument(
+                    "Output array %d shard %d is assigned more than once",
+                    mapping.out_array, out_shard);
+              }
+              if (in_buffers[in_shard] == nullptr) {
+                return InvalidArgument(
+                    "Input array %d shard %d is used more than once",
+                    mapping.in_array, in_shard);
+              }
+              out_buffers[out_shard] = std::move(in_buffers[in_shard]);
+              in_shard += in_interval.step;
+              out_shard += out_interval.step;
+            }
+            break;
+          default:
+            return InvalidArgument("Invalid ArrayCopySemantics: %d", semantics);
+        }
+      }
+    }
+  }
+
+  for (int i = 0; i < num_outputs; ++i) {
+    for (int s = 0; s < out_buffers_list[i].size(); ++s) {
+      if (out_buffers_list[i][s] == nullptr) {
+        return InvalidArgument(
+            "Output array %d addressable shard %d is unassigned", i, s);
       }
     }
   }
diff --git a/third_party/xla/xla/service/gpu/triton_fusion_analysis_test.cc b/third_party/xla/xla/service/gpu/triton_fusion_analysis_test.cc
index cde00093180611..74325b23ede671 100644
--- a/third_party/xla/xla/service/gpu/triton_fusion_analysis_test.cc
+++ b/third_party/xla/xla/service/gpu/triton_fusion_analysis_test.cc
@@ -42,8 +42,8 @@ using ::testing::FieldsAre;
 using TritonDotAnalysisTest = HloHardwareIndependentTestBase;
 
 TEST_F(TritonDotAnalysisTest, QueryingOutputScopeParametersAlwaysWorks) {
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
-                          ParseAndReturnVerifiedModule(R"(
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
+                       ParseAndReturnVerifiedModule(R"(
 triton_dot {
   p0 = f32[8,8] parameter(0)
   ROOT dot = f32[8,8] dot(p0, p0),
@@ -54,7 +54,7 @@ ENTRY e {
   p0 = f32[8,8] parameter(0)
   ROOT r = f32[8,8] fusion(p0), kind=kCustom, calls=triton_dot
 })"));
-  TF_ASSERT_OK_AND_ASSIGN(
+  ASSERT_OK_AND_ASSIGN(
       const auto analysis,
       TritonFusionAnalysis::Execute(*module->entry_computation()
                                          ->root_instruction()
@@ -85,16 +85,16 @@ ENTRY e {
     called_computations={triton_dot}
   ROOT bitcast.2 = bf16[1,8,6,3]{3,2,1,0} bitcast(custom-call)
 })";
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
-                          ParseAndReturnVerifiedModule(hlo_text));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
+                       ParseAndReturnVerifiedModule(hlo_text));
   const HloComputation* dot_computation = module->entry_computation()
                                               ->root_instruction()
                                               ->operand(0)
                                               ->called_computations()[0];
   const HloInstruction* p0 = dot_computation->parameter_instruction(0);
   const HloInstruction* p1 = dot_computation->parameter_instruction(1);
-  TF_ASSERT_OK_AND_ASSIGN(const auto analysis,
-                          TritonFusionAnalysis::Execute(*dot_computation));
+  ASSERT_OK_AND_ASSIGN(const auto analysis,
+                       TritonFusionAnalysis::Execute(*dot_computation));
   EXPECT_EQ(*analysis.ScopeParameters(TritonFusionAnalysis::Scope::LHS).begin(),
             p0);
   EXPECT_EQ(*analysis.ScopeParameters(TritonFusionAnalysis::Scope::RHS).begin(),
@@ -142,14 +142,14 @@ ENTRY e {
                          "num_stages":1,"num_warps":2,
                          "num_ctas":1}}}
 })";
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
-                          ParseAndReturnVerifiedModule(hlo_text));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
+                       ParseAndReturnVerifiedModule(hlo_text));
   const HloComputation* dot_computation =
       module->entry_computation()->root_instruction()->called_computations()[0];
   const HloInstruction* p0 = dot_computation->parameter_instruction(0);
   const HloInstruction* p1 = dot_computation->parameter_instruction(1);
-  TF_ASSERT_OK_AND_ASSIGN(const auto analysis,
-                          TritonFusionAnalysis::Execute(*dot_computation));
+  ASSERT_OK_AND_ASSIGN(const auto analysis,
+                       TritonFusionAnalysis::Execute(*dot_computation));
   EXPECT_EQ(*analysis.ScopeParameters(TritonFusionAnalysis::Scope::LHS).begin(),
             p0);
   EXPECT_EQ(*analysis.ScopeParameters(TritonFusionAnalysis::Scope::RHS).begin(),
@@ -197,16 +197,16 @@ ENTRY e {
     called_computations={triton_dot}
   ROOT bitcast.2 = bf16[1,8,6,3]{3,2,1,0} bitcast(custom-call)
 })";
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
-                          ParseAndReturnVerifiedModule(hlo_text));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
+                       ParseAndReturnVerifiedModule(hlo_text));
   const HloComputation* dot_computation = module->entry_computation()
                                               ->root_instruction()
                                               ->operand(0)
                                               ->called_computations()[0];
   const HloInstruction* p0 = dot_computation->parameter_instruction(0);
   const HloInstruction* p1 = dot_computation->parameter_instruction(1);
-  TF_ASSERT_OK_AND_ASSIGN(const auto analysis,
-                          TritonFusionAnalysis::Execute(*dot_computation));
+  ASSERT_OK_AND_ASSIGN(const auto analysis,
+                       TritonFusionAnalysis::Execute(*dot_computation));
   EXPECT_EQ(*analysis.ScopeParameters(TritonFusionAnalysis::Scope::LHS).begin(),
             p0);
   EXPECT_EQ(*analysis.ScopeParameters(TritonFusionAnalysis::Scope::RHS).begin(),
@@ -253,14 +253,14 @@ ENTRY e {
     custom_call_target="__triton",
     called_computations={triton_dot}
 })";
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
-                          ParseAndReturnVerifiedModule(hlo_text));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
+                       ParseAndReturnVerifiedModule(hlo_text));
   const HloComputation* dot_computation =
       module->entry_computation()->root_instruction()->called_computations()[0];
   const HloInstruction* p0 = dot_computation->parameter_instruction(0);
   const HloInstruction* p1 = dot_computation->parameter_instruction(1);
-  TF_ASSERT_OK_AND_ASSIGN(const auto analysis,
-                          TritonFusionAnalysis::Execute(*dot_computation));
+  ASSERT_OK_AND_ASSIGN(const auto analysis,
+                       TritonFusionAnalysis::Execute(*dot_computation));
   EXPECT_EQ(*analysis.ScopeParameters(TritonFusionAnalysis::Scope::LHS).begin(),
             p1);
   EXPECT_EQ(*analysis.ScopeParameters(TritonFusionAnalysis::Scope::RHS).begin(),
@@ -309,16 +309,16 @@ ENTRY e {
     called_computations={triton_dot}
   ROOT bitcast.2 = bf16[1,8,6,3]{3,2,1,0} bitcast(custom-call)
 })";
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
-                          ParseAndReturnVerifiedModule(hlo_text));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
+                       ParseAndReturnVerifiedModule(hlo_text));
   const HloComputation* dot_computation = module->entry_computation()
                                               ->root_instruction()
                                               ->operand(0)
                                               ->called_computations()[0];
   const HloInstruction* p0 = dot_computation->parameter_instruction(0);
   const HloInstruction* p1 = dot_computation->parameter_instruction(1);
-  TF_ASSERT_OK_AND_ASSIGN(const auto analysis,
-                          TritonFusionAnalysis::Execute(*dot_computation));
+  ASSERT_OK_AND_ASSIGN(const auto analysis,
+                       TritonFusionAnalysis::Execute(*dot_computation));
   EXPECT_EQ(*analysis.ScopeParameters(TritonFusionAnalysis::Scope::LHS).begin(),
             p0);
   EXPECT_EQ(*analysis.ScopeParameters(TritonFusionAnalysis::Scope::RHS).begin(),
@@ -368,16 +368,16 @@ ENTRY e {
     called_computations={triton_dot}
   ROOT bitcast.2 = bf16[1,8,6,3]{3,2,1,0} bitcast(custom-call)
 })";
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
-                          ParseAndReturnVerifiedModule(hlo_text));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
+                       ParseAndReturnVerifiedModule(hlo_text));
   const HloComputation* dot_computation = module->entry_computation()
                                               ->root_instruction()
                                               ->operand(0)
                                               ->called_computations()[0];
   const HloInstruction* p0 = dot_computation->parameter_instruction(0);
   const HloInstruction* p1 = dot_computation->parameter_instruction(1);
-  TF_ASSERT_OK_AND_ASSIGN(const auto analysis,
-                          TritonFusionAnalysis::Execute(*dot_computation));
+  ASSERT_OK_AND_ASSIGN(const auto analysis,
+                       TritonFusionAnalysis::Execute(*dot_computation));
   EXPECT_EQ(*analysis.ScopeParameters(TritonFusionAnalysis::Scope::LHS).begin(),
             p0);
   EXPECT_EQ(*analysis.ScopeParameters(TritonFusionAnalysis::Scope::RHS).begin(),
@@ -424,16 +424,16 @@ ENTRY e {
     custom_call_target="__triton", called_computations={triton_dot}
   ROOT bitcast.2 = bf16[3,8,1,3]{3,2,1,0} bitcast(custom-call)
 })";
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
-                          ParseAndReturnVerifiedModule(hlo_text));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
+                       ParseAndReturnVerifiedModule(hlo_text));
   const HloComputation* dot_computation = module->entry_computation()
                                               ->root_instruction()
                                               ->operand(0)
                                               ->called_computations()[0];
   const HloInstruction* p0 = dot_computation->parameter_instruction(0);
   const HloInstruction* p1 = dot_computation->parameter_instruction(1);
-  TF_ASSERT_OK_AND_ASSIGN(const auto analysis,
-                          TritonFusionAnalysis::Execute(*dot_computation));
+  ASSERT_OK_AND_ASSIGN(const auto analysis,
+                       TritonFusionAnalysis::Execute(*dot_computation));
   EXPECT_EQ(*analysis.ScopeParameters(TritonFusionAnalysis::Scope::LHS).begin(),
             p0);
   EXPECT_EQ(*analysis.ScopeParameters(TritonFusionAnalysis::Scope::RHS).begin(),
@@ -483,13 +483,13 @@ ENTRY e {
   ROOT r = bf16[3,12,2]{2,1,0} fusion(p0, p1), kind=kCustom,
     calls=triton_dot
 })";
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
-                          ParseAndReturnVerifiedModule(hlo_text));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
+                       ParseAndReturnVerifiedModule(hlo_text));
   const HloComputation* dot_computation =
       module->entry_computation()->root_instruction()->called_computations()[0];
   const HloInstruction* dot_output = dot_computation->root_instruction();
-  TF_ASSERT_OK_AND_ASSIGN(const auto analysis,
-                          TritonFusionAnalysis::Execute(*dot_computation));
+  ASSERT_OK_AND_ASSIGN(const auto analysis,
+                       TritonFusionAnalysis::Execute(*dot_computation));
   EXPECT_THAT(
       *analysis.IterSpec(TritonFusionAnalysis::Scope::OUTPUT, dot_output, 0),
       ElementsAre(FieldsAre(/*stride=*/1, /*count=*/24, /*slice_start=*/0,
@@ -505,8 +505,8 @@ ENTRY e {
 }
 
 TEST_F(TritonDotAnalysisTest, OutputParameterIsHandled) {
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
-                          ParseAndReturnVerifiedModule(R"(
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
+                       ParseAndReturnVerifiedModule(R"(
 HloModule t
 
 triton_dot {
@@ -531,8 +531,8 @@ ENTRY e {
       module->entry_computation()->root_instruction()->called_computations()[0];
   const HloInstruction* output_param =
       dot_computation->parameter_instruction(2);
-  TF_ASSERT_OK_AND_ASSIGN(const auto analysis,
-                          TritonFusionAnalysis::Execute(*dot_computation));
+  ASSERT_OK_AND_ASSIGN(const auto analysis,
+                       TritonFusionAnalysis::Execute(*dot_computation));
   EXPECT_EQ(
       analysis.IterSpec(TritonFusionAnalysis::Scope::OUTPUT, output_param, 0)
           ->size(),
@@ -556,8 +556,8 @@ ENTRY e {
 }
 
 TEST_F(TritonDotAnalysisTest, InputBroadcastFromScalarIsHandled) {
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
-                          ParseAndReturnVerifiedModule(R"(
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
+                       ParseAndReturnVerifiedModule(R"(
 HloModule t
 
 triton_dot {
@@ -577,8 +577,8 @@ ENTRY e {
   const HloComputation* dot_computation =
       module->entry_computation()->root_instruction()->called_computations()[0];
   const HloInstruction* scalar = dot_computation->parameter_instruction(1);
-  TF_ASSERT_OK_AND_ASSIGN(const auto analysis,
-                          TritonFusionAnalysis::Execute(*dot_computation));
+  ASSERT_OK_AND_ASSIGN(const auto analysis,
+                       TritonFusionAnalysis::Execute(*dot_computation));
   EXPECT_EQ(analysis.IterSpec(TritonFusionAnalysis::Scope::RHS, scalar, 0),
             nullptr);
   EXPECT_EQ(analysis.IterSpec(TritonFusionAnalysis::Scope::RHS, scalar, 1),
@@ -586,8 +586,8 @@ ENTRY e {
 }
 
 TEST_F(TritonDotAnalysisTest, InputBroadcastFromVectorIsHandled) {
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
-                          ParseAndReturnVerifiedModule(R"(
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
+                       ParseAndReturnVerifiedModule(R"(
 HloModule t
 
 triton_dot {
@@ -607,8 +607,8 @@ ENTRY e {
   const HloComputation* dot_computation =
       module->entry_computation()->root_instruction()->called_computations()[0];
   const HloInstruction* vector = dot_computation->parameter_instruction(1);
-  TF_ASSERT_OK_AND_ASSIGN(const auto analysis,
-                          TritonFusionAnalysis::Execute(*dot_computation));
+  ASSERT_OK_AND_ASSIGN(const auto analysis,
+                       TritonFusionAnalysis::Execute(*dot_computation));
   EXPECT_EQ(
       analysis.IterSpec(TritonFusionAnalysis::Scope::RHS, vector, 0)->size(),
       1);
@@ -620,8 +620,8 @@ ENTRY e {
 }
 
 TEST_F(TritonDotAnalysisTest, BroadcastFromTriviallySizedDimensionIsSupported) {
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
-                          ParseAndReturnVerifiedModule(R"(
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
+                       ParseAndReturnVerifiedModule(R"(
 f {
   p0 = f16[2] parameter(0)
   bc0 = f16[1,2] bitcast(p0)
@@ -641,8 +641,8 @@ e {
                                                ->root_instruction()
                                                ->called_computations()[0];
   const HloInstruction* p0 = dot_computation.parameter_instruction(0);
-  TF_ASSERT_OK_AND_ASSIGN(const auto analysis,
-                          TritonFusionAnalysis::Execute(dot_computation));
+  ASSERT_OK_AND_ASSIGN(const auto analysis,
+                       TritonFusionAnalysis::Execute(dot_computation));
   EXPECT_EQ(analysis.IterSpec(TritonFusionAnalysis::Scope::LHS, p0, 0)->size(),
             1);
   EXPECT_THAT(*analysis.IterSpec(TritonFusionAnalysis::Scope::LHS, p0, 0),
@@ -658,8 +658,8 @@ e {
 }
 
 TEST_F(TritonDotAnalysisTest, BroadcastWithinDimensionIsNotSupported) {
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
-                          ParseAndReturnVerifiedModule(R"(
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
+                       ParseAndReturnVerifiedModule(R"(
 f {
   a = f16[5,7] parameter(0)
   br = f16[2,5,7] broadcast(a), dimensions={1,2}
@@ -684,8 +684,8 @@ e {
 }
 
 TEST_F(TritonDotAnalysisTest, OutputBroadcastIsNotAccepted) {
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
-                          ParseAndReturnVerifiedModule(R"(
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
+                       ParseAndReturnVerifiedModule(R"(
 HloModule t
 
 ENTRY e {
@@ -705,8 +705,8 @@ ENTRY e {
 }
 
 TEST_F(TritonDotAnalysisTest, DegenerateSplitFragmentIsHandled) {
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
-                          ParseAndReturnVerifiedModule(R"(
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
+                       ParseAndReturnVerifiedModule(R"(
 triton_gemm_r {
   Arg_0.1 = s8[30,913,8,21]{3,2,1,0} parameter(0)
   bitcast.6 = s8[30,8,21,913]{2,1,3,0} bitcast(Arg_0.1)
@@ -730,8 +730,8 @@ ENTRY e {
 })"));
   const HloComputation* dot_computation =
       module->entry_computation()->root_instruction()->called_computations()[0];
-  TF_ASSERT_OK_AND_ASSIGN(const auto analysis,
-                          TritonFusionAnalysis::Execute(*dot_computation));
+  ASSERT_OK_AND_ASSIGN(const auto analysis,
+                       TritonFusionAnalysis::Execute(*dot_computation));
   EXPECT_THAT(*analysis.IterSpec(TritonFusionAnalysis::Scope::OUTPUT,
                                  dot_computation->root_instruction(), 0),
               ElementsAre(FieldsAre(/*stride=*/1, /*count=*/8 * 21,
@@ -748,8 +748,8 @@ TEST_F(TritonDotAnalysisTest,
        HandlesFurtherPropagationFromTrivialSizedTensorGracefully) {
   // We could probably support this better, just checking to avoid a crash for
   // now.
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
-                          ParseAndReturnVerifiedModule(R"(
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
+                       ParseAndReturnVerifiedModule(R"(
 triton_gemm_r {
   a = f32[3,3]{1,0} parameter(0)
   constant = f32[1,1]{1,0} constant({ {0} })
@@ -777,8 +777,8 @@ ENTRY e {
 }
 
 TEST_F(TritonDotAnalysisTest, DynamicSliceIsSupported) {
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
-                          ParseAndReturnVerifiedModule(R"(
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
+                       ParseAndReturnVerifiedModule(R"(
 triton_gemm {
   dot_lhs = f32[2,18]{1,0} parameter(0)
   dynamic_slice_input = f32[96,2]{1,0} parameter(1)
@@ -806,8 +806,8 @@ ENTRY e {
 
   const HloComputation* dot_computation =
       module->entry_computation()->root_instruction()->called_computations()[0];
-  TF_ASSERT_OK_AND_ASSIGN(const auto analysis,
-                          TritonFusionAnalysis::Execute(*dot_computation));
+  ASSERT_OK_AND_ASSIGN(const auto analysis,
+                       TritonFusionAnalysis::Execute(*dot_computation));
   const HloInstruction* p0 = dot_computation->parameter_instruction(0);
   const HloInstruction* p1 = dot_computation->parameter_instruction(1);
   EXPECT_EQ(*analysis.ScopeParameters(TritonFusionAnalysis::Scope::LHS).begin(),
@@ -837,8 +837,8 @@ ENTRY e {
 }
 
 TEST_F(TritonDotAnalysisTest, QueryScopeAlwaysWorks) {
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
-                          ParseAndReturnVerifiedModule(R"(
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
+                       ParseAndReturnVerifiedModule(R"(
 triton_gemm_r {
   Arg_0.1 = s8[30,913,8,21]{3,2,1,0} parameter(0)
   bitcast.6 = s8[30,8,21,913]{2,1,3,0} bitcast(Arg_0.1)
@@ -862,8 +862,8 @@ ENTRY e {
 })"));
   const HloComputation* dot_computation =
       module->entry_computation()->root_instruction()->called_computations()[0];
-  TF_ASSERT_OK_AND_ASSIGN(const auto analysis,
-                          TritonFusionAnalysis::Execute(*dot_computation));
+  ASSERT_OK_AND_ASSIGN(const auto analysis,
+                       TritonFusionAnalysis::Execute(*dot_computation));
   for (const auto& hlo : dot_computation->instructions()) {
     if (hlo->opcode() != HloOpcode::kDot) {
       EXPECT_TRUE(analysis.QueryInstructionScope(*hlo).has_value());
@@ -872,8 +872,8 @@ ENTRY e {
 }
 
 TEST_F(TritonDotAnalysisTest, ScaledDotIsSupported) {
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
-                          ParseAndReturnVerifiedModule(R"(
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
+                       ParseAndReturnVerifiedModule(R"(
 HloModule test
 
 scaled_dot {
@@ -886,8 +886,8 @@ scaled_dot {
       rhs_batch_dims={0}, rhs_contracting_dims={1}
 })"));
   const HloComputation* dot_computation = *module->computations().begin();
-  TF_ASSERT_OK_AND_ASSIGN(const auto analysis,
-                          TritonFusionAnalysis::Execute(*dot_computation));
+  ASSERT_OK_AND_ASSIGN(const auto analysis,
+                       TritonFusionAnalysis::Execute(*dot_computation));
   const HloInstruction* lhs = dot_computation->parameter_instruction(0);
   const HloInstruction* rhs = dot_computation->parameter_instruction(1);
   const HloInstruction* lhs_scale = dot_computation->parameter_instruction(2);
diff --git a/third_party/xla/xla/tests/constraint_propagator.cc b/third_party/xla/xla/tests/constraint_propagator.cc
index 15399c6d887e8f..70b33e7bd1000a 100644
--- a/third_party/xla/xla/tests/constraint_propagator.cc
+++ b/third_party/xla/xla/tests/constraint_propagator.cc
@@ -706,6 +706,15 @@ absl::Status ConstraintPropagator::SeedConstraints(
             states_[fused_computation->parameter_instruction(i)];
         states_[inst->operand(i)].Merge(source_state);
       }
+    } else if (inst->opcode() == HloOpcode::kCall) {
+      const HloComputation* called_computation = inst->to_apply();
+      if (called_computation != nullptr) {
+        for (int i = 0; i < inst->operand_count(); ++i) {
+          ConstraintState source_state =
+              states_[called_computation->parameter_instruction(i)];
+          states_[inst->operand(i)].Merge(source_state);
+        }
+      }
     }
   }
   return absl::OkStatus();
@@ -865,7 +874,11 @@ absl::Status ConstraintPropagator::SeedMLPatternsConstraints(
 absl::Status ConstraintPropagator::PropagateConstraintsExact(
     const HloInstruction* instruction) {
   if (instruction->opcode() == HloOpcode::kFusion) {
-    return PropagateFusionBoundary(instruction);
+    return PropagateComputationBoundary(
+        instruction, instruction->fused_instructions_computation());
+  }
+  if (instruction->opcode() == HloOpcode::kCall) {
+    return PropagateComputationBoundary(instruction, instruction->to_apply());
   }
 
   ConstraintState output_state = states_[instruction];
@@ -989,41 +1002,40 @@ absl::Status ConstraintPropagator::PropagateConstraintsExact(
   return absl::OkStatus();
 }
 
-absl::Status ConstraintPropagator::PropagateFusionBoundary(
-    const HloInstruction* fusion_instruction) {
-  const HloComputation* fused_comp =
-      fusion_instruction->fused_instructions_computation();
-  if (fused_comp == nullptr) {
+absl::Status ConstraintPropagator::PropagateComputationBoundary(
+    const HloInstruction* caller_instruction,
+    const HloComputation* callee_computation) {
+  if (callee_computation == nullptr) {
     return absl::OkStatus();
   }
 
   // 1. Output / Root binding:
-  const HloInstruction* fused_root =
-      fusion_instruction->fused_expression_root();
-  if (fused_root != nullptr) {
-    ConstraintState fusion_state = states_[fusion_instruction];
-    ConstraintState root_state = states_[fused_root];
-    // Backward: outer constraint on fusion result flows into inner root.
-    states_[fused_root].Merge(fusion_state);
-    // Forward: internal constraint computed on root flows out to fusion result.
-    states_[fusion_instruction].Merge(root_state);
+  const HloInstruction* callee_root = callee_computation->root_instruction();
+  if (callee_root != nullptr) {
+    ConstraintState caller_state = states_[caller_instruction];
+    ConstraintState root_state = states_[callee_root];
+    // Backward: outer constraint on caller result flows into inner root.
+    states_[callee_root].Merge(caller_state);
+    // Forward: internal constraint computed on root flows out to caller result.
+    states_[caller_instruction].Merge(root_state);
   }
 
   // 2. Operands / Parameters binding:
-  for (int64_t i = 0; i < fusion_instruction->operand_count(); ++i) {
-    const HloInstruction* operand = fusion_instruction->operand(i);
-    const HloInstruction* fused_param = fusion_instruction->fused_parameter(i);
-    if (fused_param == nullptr) {
+  for (int64_t i = 0; i < caller_instruction->operand_count(); ++i) {
+    const HloInstruction* operand = caller_instruction->operand(i);
+    const HloInstruction* callee_param =
+        callee_computation->parameter_instruction(i);
+    if (callee_param == nullptr) {
       continue;
     }
     ConstraintState operand_state = states_[operand];
-    ConstraintState param_state = states_[fused_param];
+    ConstraintState param_state = states_[callee_param];
     // Backward: constraints accumulated on the internal parameter flow out
     // to the caller operand.
     states_[operand].Merge(param_state);
     // Forward: constraints established on the caller operand flow into the
     // internal parameter.
-    states_[fused_param].Merge(operand_state);
+    states_[callee_param].Merge(operand_state);
   }
 
   return absl::OkStatus();
diff --git a/third_party/xla/xla/tests/constraint_propagator.h b/third_party/xla/xla/tests/constraint_propagator.h
index d70045c768df57..f78f34369598c5 100644
--- a/third_party/xla/xla/tests/constraint_propagator.h
+++ b/third_party/xla/xla/tests/constraint_propagator.h
@@ -108,11 +108,12 @@ class ConstraintPropagator {
   // formatting can simply propagate the exact constraints to their operands.
   absl::Status PropagateConstraintsExact(const HloInstruction* instruction);
 
-  // Propagates constraints bidirectionally across a kFusion instruction
-  // boundary, mapping caller operands to callee fused parameters and caller
-  // result to the fused expression root.
-  absl::Status PropagateFusionBoundary(
-      const HloInstruction* fusion_instruction);
+  // Propagates constraints bidirectionally across a subcomputation boundary
+  // (e.g. kFusion or kCall), mapping caller operands to callee parameters and
+  // caller result to the callee root instruction.
+  absl::Status PropagateComputationBoundary(
+      const HloInstruction* caller_instruction,
+      const HloComputation* callee_computation);
 
   // Propagates constraints from the output of an instruction to its operands.
   // This is approximate and introduces approximations for ops like add, sub,
diff --git a/third_party/xla/xla/tests/constraint_propagator_test.cc b/third_party/xla/xla/tests/constraint_propagator_test.cc
index 340d0f881b65aa..04c53c8656f156 100644
--- a/third_party/xla/xla/tests/constraint_propagator_test.cc
+++ b/third_party/xla/xla/tests/constraint_propagator_test.cc
@@ -1365,5 +1365,34 @@ ENTRY main {
   EXPECT_TRUE(p0_int.IsPositive());
   EXPECT_GE(p0_int.min, 0.0);
 }
+
+TEST_F(ConstraintPropagatorTest, CallPropagates) {
+  const char* hlo = R"(
+HloModule TestModule
+
+%sc_gather_comp (operand: f32[32,128], indices: s32[256]) -> f32[256,128] {
+  %operand = f32[32,128] parameter(0)
+  %indices = s32[256] parameter(1)
+  ROOT %gather = f32[256,128] gather(%operand, %indices),
+    offset_dims={1}, collapsed_slice_dims={0}, start_index_map={0},
+    index_vector_dim=1, slice_sizes={1,128}
+}
+
+ENTRY main {
+  %operand = f32[32,128] parameter(0)
+  %indices = s32[256] parameter(1)
+  ROOT %call = f32[256,128] call(%operand, %indices), to_apply=%sc_gather_comp
+}
+)";
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo));
+  ASSERT_OK_AND_ASSIGN(auto states, ConstraintPropagator::Run(*module));
+
+  auto indices_int =
+      states[module->entry_computation()->parameter_instruction(1)]
+          .GetConstraintInterval();
+  EXPECT_FALSE(indices_int.IsEmpty());
+  EXPECT_DOUBLE_EQ(indices_int.min, 0.0);
+  EXPECT_DOUBLE_EQ(indices_int.max, 31.0);
+}
 }  // namespace
 }  // namespace xla