diff --git a/tensorflow/compiler/mlir/tensorflow/utils/translate_utils.cc b/tensorflow/compiler/mlir/tensorflow/utils/translate_utils.cc
index f8277d85ab3068..ee6d90d63f3bc2 100644
--- a/tensorflow/compiler/mlir/tensorflow/utils/translate_utils.cc
+++ b/tensorflow/compiler/mlir/tensorflow/utils/translate_utils.cc
@@ -120,9 +120,8 @@ namespace {
 // Sets type list attribute with the given `name` to the given `types`. If the
 // attribute already exists with a different value, returns an error.
 template <typename ContainerT,
-          typename = typename std::enable_if<
-              std::is_same<mlir::Type, decltype(*std::declval<ContainerT>()
-                                                     .begin())>::value>::type>
+          typename = std::enable_if_t<std::is_same_v<
+              mlir::Type, decltype(*std::declval<ContainerT>().begin())>>>
 absl::Status SetTypeAttribute(absl::string_view name, ContainerT types,
                               AttrValueMap* values) {
   AttrValue value;
@@ -143,9 +142,9 @@ absl::Status SetTypeAttribute(absl::string_view name, ContainerT types,
 // Sets shape list attribute with the given `name` to the given `shapes`. If the
 // attribute already exists then this will just retain the set value.
 template <typename ContainerT,
-          typename = typename std::enable_if<std::is_same<
-              std::optional<llvm::ArrayRef<int64_t>>,
-              decltype(*std::declval<ContainerT>().begin())>::value>::type>
+          typename = std::enable_if_t<
+              std::is_same_v<std::optional<llvm::ArrayRef<int64_t>>,
+                             decltype(*std::declval<ContainerT>().begin())>>>
 void SetShapeAttribute(absl::string_view name, ContainerT shapes,
                        AttrValueMap* values) {
   AttrValue value;
diff --git a/tensorflow/core/common_runtime/gpu/BUILD b/tensorflow/core/common_runtime/gpu/BUILD
index 65d3944c9c1b43..fbfb30d068fe55 100644
--- a/tensorflow/core/common_runtime/gpu/BUILD
+++ b/tensorflow/core/common_runtime/gpu/BUILD
@@ -327,7 +327,7 @@ tf_cc_test(
 
 tf_cuda_cc_test(
     name = "gpu_bfc_allocator_test",
-    size = "small",
+    size = "medium",
     srcs = [
         "gpu_bfc_allocator_test.cc",
     ],
diff --git a/tensorflow/core/grappler/optimizers/BUILD b/tensorflow/core/grappler/optimizers/BUILD
index 00c26c46dbdd2a..8e059351395b8d 100644
--- a/tensorflow/core/grappler/optimizers/BUILD
+++ b/tensorflow/core/grappler/optimizers/BUILD
@@ -320,7 +320,7 @@ cc_library(
 
 tf_cuda_cc_test(
     name = "arithmetic_optimizer_test",
-    size = "small",
+    size = "medium",
     srcs = ["arithmetic_optimizer_test.cc"],
     deps = [
         ":arithmetic_optimizer",
diff --git a/tensorflow/python/_pywrap_tensorflow.def b/tensorflow/python/_pywrap_tensorflow.def
index b4d1cda8e1a811..9241c969219bdf 100644
--- a/tensorflow/python/_pywrap_tensorflow.def
+++ b/tensorflow/python/_pywrap_tensorflow.def
@@ -92,6 +92,7 @@ EXPORTS
  ??0MetaGraphDef@tensorflow@@IEAA@PEAVArena@protobuf@google@@@Z
  ??0MetaOptimizer@grappler@tensorflow@@QEAA@PEAVDeviceBase@2@AEBVConfigProto@2@@Z
  ??0ModelAnalyzer@grappler@tensorflow@@QEAA@AEBUGrapplerItem@12@@Z
+ ??0ModelProto_Node@model@data@tensorflow@@IEAA@PEAVArena@protobuf@google@@AEBV0123@@Z
  ??0Name@register_kernel@tensorflow@@QEAA@PEBD@Z
  ??0NamedDevice@tensorflow@@IEAA@PEAVArena@protobuf@google@@@Z
  ??0NamedDevice@tensorflow@@IEAA@PEAVArena@protobuf@google@@AEBV01@@Z
@@ -163,6 +164,7 @@ EXPORTS
  ??1MeshProto@dtensor@tensorflow@@UEAA@XZ
  ??1MessageDifferencer@util@protobuf@google@@QEAA@XZ
  ??1MetaGraphDef@tensorflow@@UEAA@XZ
+ ??1ModelProto_Node@model@data@tensorflow@@UEAA@XZ
  ??1NamedDevice@tensorflow@@UEAA@XZ
  ??1NodeDef@tensorflow@@UEAA@XZ
  ??1OpDef@tensorflow@@UEAA@XZ
diff --git a/tensorflow/python/kernel_tests/linalg/matrix_exponential_op_test.py b/tensorflow/python/kernel_tests/linalg/matrix_exponential_op_test.py
index 95b11e6b3aef12..534ed120638b32 100644
--- a/tensorflow/python/kernel_tests/linalg/matrix_exponential_op_test.py
+++ b/tensorflow/python/kernel_tests/linalg/matrix_exponential_op_test.py
@@ -19,11 +19,15 @@
 import numpy as np
 
 from tensorflow.python.client import session
+from tensorflow.python.eager import backprop
+from tensorflow.python.eager import forwardprop
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import test_util
 from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import array_ops_stack
 from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import random_ops
 from tensorflow.python.ops import variables
 from tensorflow.python.ops.linalg import linalg_impl
@@ -162,6 +166,75 @@ def testConcurrentExecutesWithoutError(self):
       expm = self.evaluate([expm1, expm2])
       self.assertAllEqual(expm[0], expm[1])
 
+  @test_util.run_v2_only
+  def testSecondDerivativeNestedForwardAccumulator(self):
+    # Regression test for Issue #127224.
+    def target(t):
+      one = constant_op.constant(1.0, dtype=t.dtype)
+      a = array_ops_stack.stack([
+          array_ops_stack.stack([t, one]),
+          array_ops_stack.stack([-one, t]),
+      ])
+      return math_ops.reduce_sum(linalg_impl.matrix_exponential(a))
+
+    def forward_ad(fn):
+      def jvp(x):
+        with forwardprop.ForwardAccumulator(x, array_ops.ones_like(x)) as acc:
+          y = fn(x)
+        return acc.jvp(y)
+
+      return jvp
+
+    for dtype in [ops.dtypes.float32, ops.dtypes.float64]:
+      x = constant_op.constant(0.7, dtype=dtype)
+      expected = 2.0 * np.exp(0.7) * np.cos(1.0)
+
+      if dtype == ops.dtypes.float64:
+        atol_1st = 1e-12
+        rtol_1st = 1e-12
+        # Nested ForwardAccumulator (forward-over-forward) diverges from the
+        # analytical value by ~9% here, independent of dtype: matrix_exponential
+        # chains matmuls three levels deep (matrix_4 = matrix_2 * matrix_2,
+        # matrix_6 = matrix_4 * matrix_2, ...), and TF's forward-over-forward
+        # differentiation is inaccurate past two chained matmul levels
+        # regardless of precision, so this isn't something this PR's
+        # coefficient normalization can improve further.
+        atol_2nd = 0.25
+        rtol_2nd = 1e-1
+        atol_hess = 1e-12
+        rtol_hess = 1e-12
+      else:
+        atol_1st = 1e-3
+        rtol_1st = 1e-3
+        atol_2nd = 0.25
+        rtol_2nd = 1e-1
+        atol_hess = 1e-4
+        rtol_hess = 1e-4
+
+      # 1st order forward-mode JVP
+      actual_1st = forward_ad(target)(x)
+      self.assertAllClose(
+          expected, self.evaluate(actual_1st), rtol=rtol_1st, atol=atol_1st
+      )
+
+      # 2nd order nested forward-mode JVP
+      actual_2nd = forward_ad(forward_ad(target))(x)
+      self.assertAllClose(
+          expected, self.evaluate(actual_2nd), rtol=rtol_2nd, atol=atol_2nd
+      )
+
+      # 2nd order reverse-mode Hessian
+      with backprop.GradientTape() as t2:
+        t2.watch(x)
+        with backprop.GradientTape() as t1:
+          t1.watch(x)
+          y = target(x)
+        g1 = t1.gradient(y, x)
+      g2 = t2.gradient(g1, x)
+      self.assertAllClose(
+          expected, self.evaluate(g2), rtol=rtol_hess, atol=atol_hess
+      )
+
 
 class MatrixExponentialBenchmark(test.Benchmark):
 
diff --git a/tensorflow/python/ops/linalg/linalg_impl.py b/tensorflow/python/ops/linalg/linalg_impl.py
index f6f39cf61aadad..820e33a658a025 100644
--- a/tensorflow/python/ops/linalg/linalg_impl.py
+++ b/tensorflow/python/ops/linalg/linalg_impl.py
@@ -134,39 +134,39 @@ def adjoint(matrix, name=None):
 # https://eigen.tuxfamily.org/dox/unsupported/MatrixExponential_8h_source.html
 def _matrix_exp_pade3(matrix):
   """3rd-order Pade approximant for matrix exponential."""
-  b = [120.0, 60.0, 12.0]
-  b = [constant_op.constant(x, matrix.dtype) for x in b]
+  b = [120.0, 60.0, 12.0, 1.0]
+  c = [constant_op.constant(x / b[0], matrix.dtype) for x in b]
   ident = linalg_ops.eye(
       array_ops.shape(matrix)[-2],
       batch_shape=array_ops.shape(matrix)[:-2],
       dtype=matrix.dtype)
   matrix_2 = math_ops.matmul(matrix, matrix)
-  tmp = matrix_2 + b[1] * ident
+  tmp = c[3] * matrix_2 + c[1] * ident
   matrix_u = math_ops.matmul(matrix, tmp)
-  matrix_v = b[2] * matrix_2 + b[0] * ident
+  matrix_v = c[2] * matrix_2 + ident
   return matrix_u, matrix_v
 
 
 def _matrix_exp_pade5(matrix):
   """5th-order Pade approximant for matrix exponential."""
-  b = [30240.0, 15120.0, 3360.0, 420.0, 30.0]
-  b = [constant_op.constant(x, matrix.dtype) for x in b]
+  b = [30240.0, 15120.0, 3360.0, 420.0, 30.0, 1.0]
+  c = [constant_op.constant(x / b[0], matrix.dtype) for x in b]
   ident = linalg_ops.eye(
       array_ops.shape(matrix)[-2],
       batch_shape=array_ops.shape(matrix)[:-2],
       dtype=matrix.dtype)
   matrix_2 = math_ops.matmul(matrix, matrix)
   matrix_4 = math_ops.matmul(matrix_2, matrix_2)
-  tmp = matrix_4 + b[3] * matrix_2 + b[1] * ident
+  tmp = c[5] * matrix_4 + c[3] * matrix_2 + c[1] * ident
   matrix_u = math_ops.matmul(matrix, tmp)
-  matrix_v = b[4] * matrix_4 + b[2] * matrix_2 + b[0] * ident
+  matrix_v = c[4] * matrix_4 + c[2] * matrix_2 + ident
   return matrix_u, matrix_v
 
 
 def _matrix_exp_pade7(matrix):
   """7th-order Pade approximant for matrix exponential."""
-  b = [17297280.0, 8648640.0, 1995840.0, 277200.0, 25200.0, 1512.0, 56.0]
-  b = [constant_op.constant(x, matrix.dtype) for x in b]
+  b = [17297280.0, 8648640.0, 1995840.0, 277200.0, 25200.0, 1512.0, 56.0, 1.0]
+  c = [constant_op.constant(x / b[0], matrix.dtype) for x in b]
   ident = linalg_ops.eye(
       array_ops.shape(matrix)[-2],
       batch_shape=array_ops.shape(matrix)[:-2],
@@ -174,19 +174,27 @@ def _matrix_exp_pade7(matrix):
   matrix_2 = math_ops.matmul(matrix, matrix)
   matrix_4 = math_ops.matmul(matrix_2, matrix_2)
   matrix_6 = math_ops.matmul(matrix_4, matrix_2)
-  tmp = matrix_6 + b[5] * matrix_4 + b[3] * matrix_2 + b[1] * ident
+  tmp = c[7] * matrix_6 + c[5] * matrix_4 + c[3] * matrix_2 + c[1] * ident
   matrix_u = math_ops.matmul(matrix, tmp)
-  matrix_v = b[6] * matrix_6 + b[4] * matrix_4 + b[2] * matrix_2 + b[0] * ident
+  matrix_v = c[6] * matrix_6 + c[4] * matrix_4 + c[2] * matrix_2 + ident
   return matrix_u, matrix_v
 
 
 def _matrix_exp_pade9(matrix):
   """9th-order Pade approximant for matrix exponential."""
   b = [
-      17643225600.0, 8821612800.0, 2075673600.0, 302702400.0, 30270240.0,
-      2162160.0, 110880.0, 3960.0, 90.0
+      17643225600.0,
+      8821612800.0,
+      2075673600.0,
+      302702400.0,
+      30270240.0,
+      2162160.0,
+      110880.0,
+      3960.0,
+      90.0,
+      1.0,
   ]
-  b = [constant_op.constant(x, matrix.dtype) for x in b]
+  c = [constant_op.constant(x / b[0], matrix.dtype) for x in b]
   ident = linalg_ops.eye(
       array_ops.shape(matrix)[-2],
       batch_shape=array_ops.shape(matrix)[:-2],
@@ -196,23 +204,42 @@ def _matrix_exp_pade9(matrix):
   matrix_6 = math_ops.matmul(matrix_4, matrix_2)
   matrix_8 = math_ops.matmul(matrix_6, matrix_2)
   tmp = (
-      matrix_8 + b[7] * matrix_6 + b[5] * matrix_4 + b[3] * matrix_2 +
-      b[1] * ident)
+      c[9] * matrix_8
+      + c[7] * matrix_6
+      + c[5] * matrix_4
+      + c[3] * matrix_2
+      + c[1] * ident
+  )
   matrix_u = math_ops.matmul(matrix, tmp)
   matrix_v = (
-      b[8] * matrix_8 + b[6] * matrix_6 + b[4] * matrix_4 + b[2] * matrix_2 +
-      b[0] * ident)
+      c[8] * matrix_8
+      + c[6] * matrix_6
+      + c[4] * matrix_4
+      + c[2] * matrix_2
+      + ident
+  )
   return matrix_u, matrix_v
 
 
 def _matrix_exp_pade13(matrix):
   """13th-order Pade approximant for matrix exponential."""
   b = [
-      64764752532480000.0, 32382376266240000.0, 7771770303897600.0,
-      1187353796428800.0, 129060195264000.0, 10559470521600.0, 670442572800.0,
-      33522128640.0, 1323241920.0, 40840800.0, 960960.0, 16380.0, 182.0
+      64764752532480000.0,
+      32382376266240000.0,
+      7771770303897600.0,
+      1187353796428800.0,
+      129060195264000.0,
+      10559470521600.0,
+      670442572800.0,
+      33522128640.0,
+      1323241920.0,
+      40840800.0,
+      960960.0,
+      16380.0,
+      182.0,
+      1.0,
   ]
-  b = [constant_op.constant(x, matrix.dtype) for x in b]
+  c = [constant_op.constant(x / b[0], matrix.dtype) for x in b]
   ident = linalg_ops.eye(
       array_ops.shape(matrix)[-2],
       batch_shape=array_ops.shape(matrix)[:-2],
@@ -221,13 +248,23 @@ def _matrix_exp_pade13(matrix):
   matrix_4 = math_ops.matmul(matrix_2, matrix_2)
   matrix_6 = math_ops.matmul(matrix_4, matrix_2)
   tmp_u = (
-      math_ops.matmul(matrix_6, matrix_6 + b[11] * matrix_4 + b[9] * matrix_2) +
-      b[7] * matrix_6 + b[5] * matrix_4 + b[3] * matrix_2 + b[1] * ident)
+      math_ops.matmul(
+          matrix_6, c[13] * matrix_6 + c[11] * matrix_4 + c[9] * matrix_2
+      )
+      + c[7] * matrix_6
+      + c[5] * matrix_4
+      + c[3] * matrix_2
+      + c[1] * ident
+  )
   matrix_u = math_ops.matmul(matrix, tmp_u)
-  tmp_v = b[12] * matrix_6 + b[10] * matrix_4 + b[8] * matrix_2
+  tmp_v = c[12] * matrix_6 + c[10] * matrix_4 + c[8] * matrix_2
   matrix_v = (
-      math_ops.matmul(matrix_6, tmp_v) + b[6] * matrix_6 + b[4] * matrix_4 +
-      b[2] * matrix_2 + b[0] * ident)
+      math_ops.matmul(matrix_6, tmp_v)
+      + c[6] * matrix_6
+      + c[4] * matrix_4
+      + c[2] * matrix_2
+      + ident
+  )
   return matrix_u, matrix_v
 
 
diff --git a/third_party/xla/build_tools/lint/run_dwyu.py b/third_party/xla/build_tools/lint/run_dwyu.py
index cc2174aca73bad..1fedb40d172024 100644
--- a/third_party/xla/build_tools/lint/run_dwyu.py
+++ b/third_party/xla/build_tools/lint/run_dwyu.py
@@ -32,6 +32,58 @@
 import tempfile
 from typing import Iterator, Sequence
 
+_ADD_STRING_VIEW_PREFIX = (
+    "buildozer 'add deps @com_google_absl//absl/strings:string_view' "
+)
+_REMOVE_STRINGS_PREFIXES = (
+    "buildozer 'remove deps @com_google_absl//absl/strings' ",
+    "buildozer 'remove deps @com_google_absl//absl/strings:strings' ",
+)
+_BANT_CLEANUP_FINDINGS_EXIT_CODE = 3
+
+
+def _is_paired_strings_removal(
+    line: str, string_view_targets: set[str]
+) -> bool:
+  """Return whether `line` removes `absl/strings` from a `string_view` target."""
+  for prefix in _REMOVE_STRINGS_PREFIXES:
+    if line.startswith(prefix):
+      return line.removeprefix(prefix) in string_view_targets
+  return False
+
+
+def _filter_bant_stdout(stdout: str) -> list[str]:
+  """Suppress @com_google_absl//absl/strings:string_view findings from BANT.
+
+  `@com_google_absl//absl/strings` exports `string_view.h` in `hdrs`, so Bazel's
+  `layering_check` allows including `absl/strings/string_view.h` when
+  `@com_google_absl//absl/strings` is in `deps`. However, BANT has a hardcoded
+  `absl_string_view_skip` special case that ignores `string_view.h` on
+  `absl/strings:strings`, which does not align with `layering_check`.
+
+  Args:
+    stdout: Standard output emitted by `bant dwyu`.
+
+  Returns:
+    Remaining `buildozer` finding lines after filtering.
+  """
+  lines = stdout.splitlines()
+  string_view_targets = set()
+  for line in lines:
+    if line.startswith(_ADD_STRING_VIEW_PREFIX):
+      string_view_targets.add(line.removeprefix(_ADD_STRING_VIEW_PREFIX))
+  if not string_view_targets:
+    return lines
+
+  filtered = []
+  for line in lines:
+    if line.startswith(_ADD_STRING_VIEW_PREFIX):
+      continue
+    if _is_paired_strings_removal(line, string_view_targets):
+      continue
+    filtered.append(line)
+  return filtered
+
 
 @contextlib.contextmanager
 def bant_workspace(
@@ -77,7 +129,7 @@ def run_dwyu(
       "@xla" + label if label.startswith("//") else label for label in targets
   ]
   with bant_workspace(pathlib.Path.cwd(), output_base) as workspace:
-    return subprocess.run(
+    proc = subprocess.run(
         [
             str(pathlib.Path(executable).resolve()),
             "-C",
@@ -90,8 +142,17 @@ def run_dwyu(
             "dwyu",
             *labels,
         ],
+        stdout=subprocess.PIPE,
+        text=True,
         check=False,
-    ).returncode
+    )
+  findings = _filter_bant_stdout(proc.stdout)
+  if findings:
+    sys.stdout.write("\n".join(findings) + "\n")
+    sys.stdout.flush()
+  if proc.returncode == _BANT_CLEANUP_FINDINGS_EXIT_CODE and not findings:
+    return 0
+  return proc.returncode
 
 
 def main(argv: Sequence[str]) -> int:
diff --git a/third_party/xla/build_tools/lint/run_dwyu_test.py b/third_party/xla/build_tools/lint/run_dwyu_test.py
index 16d0ba4d0c6db7..ec3d8a39cd84da 100644
--- a/third_party/xla/build_tools/lint/run_dwyu_test.py
+++ b/third_party/xla/build_tools/lint/run_dwyu_test.py
@@ -90,6 +90,33 @@ def test_empty_selection_does_not_run_bant(self):
         0,
     )
 
+  def test_filters_absl_string_view_findings(self):
+    stdout = (
+        "buildozer 'remove deps @com_google_absl//absl/strings'"
+        " @xla//xla/only_string_view\n"
+        "buildozer 'add deps @com_google_absl//absl/strings:string_view'"
+        " @xla//xla/only_string_view\n"
+        "buildozer 'add deps @com_google_absl//absl/strings:string_view'"
+        " @xla//xla/with_str_cat\n"
+        "buildozer 'remove deps @com_google_absl//absl/strings'"
+        " @xla//xla/truly_unused_strings\n"
+        "buildozer 'add deps @xla//xla/platform:errors'"
+        " @xla//xla/with_str_cat\n"
+    )
+    self.assertEqual(
+        run_dwyu._filter_bant_stdout(stdout),
+        [
+            (
+                "buildozer 'remove deps @com_google_absl//absl/strings'"
+                " @xla//xla/truly_unused_strings"
+            ),
+            (
+                "buildozer 'add deps @xla//xla/platform:errors'"
+                " @xla//xla/with_str_cat"
+            ),
+        ],
+    )
+
 
 @unittest.skipUnless(
     _BANT, "requires the BANT binary installed by the workflow"
@@ -265,6 +292,92 @@ def test_all_allowed_rules_supported_by_bant_macros(self):
         result.stderr,
     )
 
+  def _setup_absl_strings_repo(self):
+    self.write_file(
+        "MODULE.bazel",
+        """\
+module(name = "xla")
+bazel_dep(name = "abseil-cpp", version = "20260526.0", repo_name = "com_google_absl")
+""",
+    )
+    absl_repo = self.external / "abseil-cpp+"
+    (absl_repo / "absl/strings").mkdir(parents=True)
+    (absl_repo / "MODULE.bazel").write_text('module(name = "abseil-cpp")\n')
+    (absl_repo / "absl/strings/str_cat.h").write_text("// str_cat\n")
+    (absl_repo / "absl/strings/string_view.h").write_text("// string_view\n")
+    (absl_repo / "absl/strings/BUILD.bazel").write_text("""\
+cc_library(
+    name = "string_view",
+    hdrs = ["string_view.h"],
+    visibility = ["//visibility:public"],
+)
+
+cc_library(
+    name = "strings",
+    hdrs = [
+        "str_cat.h",
+        "string_view.h",
+    ],
+    textual_hdrs = [
+        "string_view.h",
+    ],
+    visibility = ["//visibility:public"],
+    deps = [":string_view"],
+)
+""")
+
+  def test_absl_strings_and_string_view_filtering(self):
+    self._setup_absl_strings_repo()
+    test_cases = [
+        (
+            "satisfies_string_view_and_str_cat_include",
+            '["@com_google_absl//absl/strings"]',
+            (
+                '#include "absl/strings/str_cat.h"\n'
+                '#include "absl/strings/string_view.h"\n'
+            ),
+            0,
+            "",
+        ),
+        (
+            "satisfies_string_view_only_include",
+            '["@com_google_absl//absl/strings"]',
+            '#include "absl/strings/string_view.h"\n',
+            0,
+            "",
+        ),
+        (
+            "unused_absl_strings_still_reported",
+            '["//xla/platform:errors", "@com_google_absl//absl/strings"]',
+            '#include "xla/platform/errors.h"\n',
+            3,
+            (
+                "buildozer 'remove deps @com_google_absl//absl/strings'"
+                " @xla//xla/consumer\n"
+            ),
+        ),
+    ]
+    for name, deps, source, expected_returncode, expected_stdout in test_cases:
+      with self.subTest(name):
+        self.write_file(
+            "xla/consumer/BUILD",
+            f"""\
+cc_library(
+    name = "consumer",
+    srcs = ["consumer.cc"],
+    deps = {deps},
+)
+""",
+        )
+        self.write_file("xla/consumer/consumer.cc", source)
+        result = self.check_targets("//xla/consumer:consumer")
+        self.assertEqual(
+            result.returncode,
+            expected_returncode,
+            result.stdout + result.stderr,
+        )
+        self.assertEqual(result.stdout, expected_stdout)
+
 
 if __name__ == "__main__":
   unittest.main()
diff --git a/third_party/xla/build_tools/rocm/rocm_xla.bazelrc b/third_party/xla/build_tools/rocm/rocm_xla.bazelrc
index df291bc229421e..546eed927c42f8 100644
--- a/third_party/xla/build_tools/rocm/rocm_xla.bazelrc
+++ b/third_party/xla/build_tools/rocm/rocm_xla.bazelrc
@@ -23,7 +23,6 @@ test:rocm_rbe --remote_timeout=3600
 test:rocm_rbe --strategy=TestRunner=remote,local
 
 # Dynamic execution config for ROCm RBE - races remote vs local with remote head start
-# Compilation, linking, and tests all race, with CppLink forced local due to toolchain issues
 build:rocm_rbe_dynamic --config=rocm_rbe
 build:rocm_rbe_dynamic --jobs=50
 build:rocm_rbe_dynamic --internal_spawn_scheduler
@@ -31,9 +30,8 @@ build:rocm_rbe_dynamic --spawn_strategy=dynamic
 build:rocm_rbe_dynamic --dynamic_local_strategy=worker,standalone,local
 build:rocm_rbe_dynamic --dynamic_remote_strategy=remote
 build:rocm_rbe_dynamic --experimental_local_execution_delay=1000
-build:rocm_rbe_dynamic --remote_download_all
+build:rocm_rbe_dynamic --remote_download_minimal
 build:rocm_rbe_dynamic --local_resources=cpu=HOST_CPUS
-build:rocm_rbe_dynamic --strategy=CppLink=standalone,local
 
 build:tsan --strip=never
 build:tsan --copt -fsanitize=thread
@@ -54,8 +52,7 @@ build:ci_multi_gpu --flaky_test_attempts=3
 build:ci_multi_gpu --experimental_guard_against_concurrent_changes
 build:ci_multi_gpu --test_env=HIP_VISIBLE_DEVICES=0,1,2,3
 
-build:ci_rocm_cpu --local_test_jobs=200
-build:ci_rocm_cpu --strategy=TestRunner=local
+build:ci_rocm_cpu --jobs=50
 
 test:xla_sgpu -- \
 //xla/... \
diff --git a/third_party/xla/tensorflow.bazelrc b/third_party/xla/tensorflow.bazelrc
index 6da958705ae12e..2b7428d955fc60 100644
--- a/third_party/xla/tensorflow.bazelrc
+++ b/third_party/xla/tensorflow.bazelrc
@@ -293,8 +293,8 @@ common:rocm_clang_hermetic --@rules_ml_toolchain//common:enable_cuda=False
 common:rocm_clang_hermetic --@rules_ml_toolchain//common:enable_sycl=False
 common:rocm_clang_hermetic --@rules_ml_toolchain//common:enable_hermetic_cc=True
 common:rocm_clang_hermetic --@local_config_rocm//rocm:rocm_path_type=hermetic
-common:rocm_clang_hermetic --strategy=CppLink=local
 common:rocm_clang_hermetic --@rules_ml_toolchain//common:static_libcxx=False
+common:rocm_clang_hermetic --strategy=CppLink=local
 
 common:rocm --config=rocm_clang_hermetic
 common:rocm_ci --config=rocm
@@ -302,6 +302,7 @@ common:rocm_ci --@local_config_rocm//rocm:rocm_path_type=hermetic
 
 common:rocm_ci_hermetic --dynamic_mode=fully
 common:rocm_ci_hermetic --@rules_ml_toolchain//common:enable_xla_test_global_symbol_version=True
+common:rocm_ci_hermetic --@rules_ml_toolchain//common:enable_hermetic_libraries=True
 common:rocm_ci_hermetic --config=rocm_clang_hermetic
 common:rocm_ci_hermetic --repo_env=TF_ROCM_AMDGPU_TARGETS="gfx908,gfx90a"
 common:rocm_ci_hermetic --repo_env=ROCM_DISTRO_VERSION="rocm_7.13.0_gfx90a"
diff --git a/third_party/xla/third_party/llvm/platforms_config.patch b/third_party/xla/third_party/llvm/platforms_config.patch
deleted file mode 100644
index 89c14990db4f41..00000000000000
--- a/third_party/xla/third_party/llvm/platforms_config.patch
+++ /dev/null
@@ -1,33 +0,0 @@
-# Copyright 2026 The OpenXLA Authors.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-diff --git a/utils/bazel/llvm-project-overlay/llvm/config.bzl b/utils/bazel/llvm-project-overlay/llvm/config.bzl
---- a/utils/bazel/llvm-project-overlay/llvm/config.bzl
-+++ b/utils/bazel/llvm-project-overlay/llvm/config.bzl
-@@ -69,14 +69,13 @@
- backtrace_defines = select({
-     "@platforms//os:emscripten": [],
-     "@platforms//os:windows": [],
--    "@rules_cc//cc/libc:musl": [],
-     "//conditions:default": [
-         "HAVE_BACKTRACE=1",
-         "BACKTRACE_HEADER=<execinfo.h>",
-     ],
- })
- 
- mallinfo_defines = select({
--    "@rules_cc//cc/libc:glibc": ["HAVE_MALLINFO=1"],
-+    "@platforms//os:linux": ["HAVE_MALLINFO=1"],
-     "//conditions:default": [],
- })
diff --git a/third_party/xla/third_party/llvm/workspace.bzl b/third_party/xla/third_party/llvm/workspace.bzl
index 62fd3f74fa8e48..83d4a833e45971 100644
--- a/third_party/xla/third_party/llvm/workspace.bzl
+++ b/third_party/xla/third_party/llvm/workspace.bzl
@@ -34,7 +34,6 @@ def repo(name):
         patch_file = [
             "//third_party/llvm:generated.patch",  # Autogenerated, don't remove.
             "//third_party/llvm:build.patch",
-            "//third_party/llvm:platforms_config.patch",
             "//third_party/llvm:toolchains.patch",
             "//third_party/llvm:lit_test.patch",
             "//third_party/llvm:run_lit.patch",
diff --git a/third_party/xla/xla/backends/cpu/codegen/tools/BUILD b/third_party/xla/xla/backends/cpu/codegen/tools/BUILD
index d9633189db5ca0..4917e39f0ec28b 100644
--- a/third_party/xla/xla/backends/cpu/codegen/tools/BUILD
+++ b/third_party/xla/xla/backends/cpu/codegen/tools/BUILD
@@ -98,7 +98,7 @@ xla_cc_binary(
 )
 
 xla_cc_binary(
-    name = "fusion_to_mlir",
+    name = "cpu_fusion_to_mlir",
     testonly = True,
     srcs = ["fusion_to_mlir.cc"],
     deps = [
diff --git a/third_party/xla/xla/backends/gpu/autotuner/factory_cuda.cc b/third_party/xla/xla/backends/gpu/autotuner/factory_cuda.cc
index ce6b1c5df295f9..8d9094a35cd758 100644
--- a/third_party/xla/xla/backends/gpu/autotuner/factory_cuda.cc
+++ b/third_party/xla/xla/backends/gpu/autotuner/factory_cuda.cc
@@ -55,17 +55,16 @@ std::unique_ptr<HloPassPipeline> GetCublasLtRewriterPipeline(
     const stream_executor::DeviceDescription& device_description) {
   auto pipeline =
       std::make_unique<HloPassPipeline>("cublaslt_rewriter_pipeline");
-  pipeline->AddPass(std::make_unique<DotAlgorithmRewriter>());
   pipeline->AddPass(std::make_unique<ScaledDotRewriter>());
-  for (GemmRewriterOptions::DType dtype :
-       {GemmRewriterOptions::DType::kFp8Only,
-        GemmRewriterOptions::DType::kNonFp8Only}) {
-    GemmRewriterOptions options{dtype};
-    auto gemm_rewriter = std::make_unique<GemmRewriter>(
-        device_description.gpu_compute_capability(),
-        device_description.runtime_version(), options);
-    pipeline->AddPass(std::move(gemm_rewriter));
-  }
+  pipeline->AddPass(std::make_unique<GemmRewriter>(
+      device_description.gpu_compute_capability(),
+      device_description.runtime_version(),
+      GemmRewriterOptions{GemmRewriterOptions::DType::kFp8Only}));
+  pipeline->AddPass(std::make_unique<DotAlgorithmRewriter>());
+  pipeline->AddPass(std::make_unique<GemmRewriter>(
+      device_description.gpu_compute_capability(),
+      device_description.runtime_version(),
+      GemmRewriterOptions{GemmRewriterOptions::DType::kNonFp8Only}));
   return pipeline;
 }
 
diff --git a/third_party/xla/xla/backends/gpu/autotuner/factory_rocm.cc b/third_party/xla/xla/backends/gpu/autotuner/factory_rocm.cc
index cc608307092991..6e195de99ff694 100644
--- a/third_party/xla/xla/backends/gpu/autotuner/factory_rocm.cc
+++ b/third_party/xla/xla/backends/gpu/autotuner/factory_rocm.cc
@@ -62,15 +62,16 @@ std::unique_ptr<HloPassPipeline> GetGemmRewriterPipeline(
     absl::Span<const DType> dtypes) {
   auto pipeline =
       std::make_unique<HloPassPipeline>("hipblaslt_rewriter_pipeline");
-  pipeline->AddPass(std::make_unique<DotAlgorithmRewriter>());
   pipeline->AddPass(std::make_unique<ScaledDotRewriter>());
-  for (DType dtype : dtypes) {
-    GemmRewriterOptions options{dtype};
-    auto gemm_rewriter = std::make_unique<GemmRewriter>(
-        device_description.gpu_compute_capability(),
-        device_description.runtime_version(), options);
-    pipeline->AddPass(std::move(gemm_rewriter));
-  }
+  pipeline->AddPass(std::make_unique<GemmRewriter>(
+      device_description.gpu_compute_capability(),
+      device_description.runtime_version(),
+      GemmRewriterOptions{DType::kFp8Only}));
+  pipeline->AddPass(std::make_unique<DotAlgorithmRewriter>());
+  pipeline->AddPass(std::make_unique<GemmRewriter>(
+      device_description.gpu_compute_capability(),
+      device_description.runtime_version(),
+      GemmRewriterOptions{DType::kNonFp8Only}));
   return pipeline;
 }
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/mlir_kernel_emitter.cc b/third_party/xla/xla/backends/gpu/codegen/emitters/mlir_kernel_emitter.cc
index a09f10b5309ba2..20ff266bc30414 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/mlir_kernel_emitter.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/mlir_kernel_emitter.cc
@@ -393,9 +393,10 @@ AsyncThunkSequence MlirKernelFusion::Emit(
                   .Map([kernel_name = std::move(kernel_name),
                         launch_dims = std::move(launch_dims),
                         use_pdl](const std::vector<uint8_t>& cubin) mutable {
-                    KernelReuseCache::Entry entry{kernel_name, launch_dims,
-                                                  std::nullopt,
-                                                  /*shmem_bytes=*/0, cubin};
+                    KernelReuseCache::Entry entry{
+                        kernel_name, launch_dims, std::nullopt,
+                        /*shmem_bytes=*/0,
+                        std::make_shared<const std::vector<uint8_t>>(cubin)};
 
                     entry.use_pdl = use_pdl;
                     return entry;
@@ -407,20 +408,20 @@ AsyncThunkSequence MlirKernelFusion::Emit(
   bool kernel_cached = cached;
   return future_entry.Map([&fusion, thunk_info = std::move(thunk_info),
                            args = std::move(args), kernel_cached](
-                              const KernelReuseCache::Entry* entry) mutable
+                              const KernelReuseCache::Entry& entry) mutable
                               -> absl::StatusOr<ThunkSequence> {
     if (kernel_cached) {
-      VLOG(3) << "Reuse: " << fusion.name() << " -> " << entry->kernel_name;
+      VLOG(3) << "Reuse: " << fusion.name() << " -> " << entry.kernel_name;
     }
     ABSL_ASSIGN_OR_RETURN(CustomKernel custom_kernel,
-                     kernel::CreateOwnedCubinCustomKernel(
-                         entry->kernel_name, entry->binary, args.args().size(),
-                         entry->launch_dimensions.block_counts(),
-                         entry->launch_dimensions.thread_counts_per_block(),
-                         entry->shmem_bytes));
+                     kernel::CreateSharedCubinCustomKernel(
+                         entry.kernel_name, entry.binary, args.args().size(),
+                         entry.launch_dimensions.block_counts(),
+                         entry.launch_dimensions.thread_counts_per_block(),
+                         entry.shmem_bytes));
 
     return ThunkSequence::Of<CustomKernelThunk>(
-        thunk_info, std::move(custom_kernel), args, entry->use_pdl);
+        thunk_info, std::move(custom_kernel), args, entry.use_pdl);
   });
 }
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/BUILD b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/BUILD
index 4201819a1b030f..c8d8e8a53b910b 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/BUILD
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/BUILD
@@ -41,7 +41,7 @@ lit_test_suite(
     default_tags = tf_cuda_tests_tags(),
     exec_properties = tf_exec_properties({"tags": tf_cuda_tests_tags()}),
     tools = [
-        "//xla/backends/gpu/codegen/tools:fusion_to_mlir",
+        "//xla/backends/gpu/codegen/tools:gpu_fusion_to_mlir",
         "//xla/backends/gpu/codegen/tools:gpu_test_correctness",
         "//xla/codegen/tools:emitters_opt",
         "@llvm-project//llvm:FileCheck",
@@ -54,7 +54,7 @@ lit_test_suite(
     cfg = "//xla:lit.cfg.py",
     default_tags = ["gpu"],
     tools = [
-        "//xla/backends/gpu/codegen/tools:fusion_to_mlir",
+        "//xla/backends/gpu/codegen/tools:gpu_fusion_to_mlir",
         "//xla/codegen/tools:emitters_opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/convert/f4e2m1fn_intrinsics.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/convert/f4e2m1fn_intrinsics.hlo
index 27ebfcf03e6e93..81e4fcac50b7f1 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/convert/f4e2m1fn_intrinsics.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/convert/f4e2m1fn_intrinsics.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -xla-gpu-convert-float-nvidia='compute_capability_major=10 compute_capability_minor=0 ptx_version_major=8 ptx_version_minor=6' | FileCheck %s
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -xla-gpu-convert-float-nvidia='compute_capability_major=10 compute_capability_minor=0 ptx_version_major=8 ptx_version_minor=6' | FileCheck %s
 // RUN: gpu_test_correctness %s
 
 f {
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_column/horizontal_fusion_two_groups.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_column/horizontal_fusion_two_groups.hlo
index f0e27a3a080904..c731b8907f76f8 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_column/horizontal_fusion_two_groups.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_column/horizontal_fusion_two_groups.hlo
@@ -1,5 +1,5 @@
 // RUN: gpu_test_correctness %s
-// RUN: fusion_to_mlir %s | FileCheck %s
+// RUN: gpu_fusion_to_mlir %s | FileCheck %s
 
 region {
   param_0 = f32[] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_column/odd_sizes.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_column/odd_sizes.hlo
index d871f840b820d8..e03ee5f3f6d22f 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_column/odd_sizes.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_column/odd_sizes.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
 // RUN:   FileCheck %s
 // RUN: gpu_test_correctness %s --bijection_inputs=reduce:0 \
 // RUN:   --bijection_outputs=reduce
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_column/vectorize_x1_c64.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_column/vectorize_x1_c64.hlo
index dbdb19b18d9c22..7ccce4dbdca1c9 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_column/vectorize_x1_c64.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_column/vectorize_x1_c64.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
 // RUN:   FileCheck %s
 // RUN: gpu_test_correctness %s --bijection_inputs=reduce:0 \
 // RUN:   --bijection_outputs=reduce
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_column/vectorize_x1_f64.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_column/vectorize_x1_f64.hlo
index 36ef46cc3797b2..e7e560766b77e0 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_column/vectorize_x1_f64.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_column/vectorize_x1_f64.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
 // RUN:   FileCheck %s
 // RUN: gpu_test_correctness %s --bijection_inputs=reduce:0 \
 // RUN:   --bijection_outputs=reduce
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_column/vectorize_x2_f32.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_column/vectorize_x2_f32.hlo
index fe182f94d6a567..776607c7cbc11b 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_column/vectorize_x2_f32.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_column/vectorize_x2_f32.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
 // RUN:   FileCheck %s
 // RUN: gpu_test_correctness %s --bijection_inputs=reduce:0 \
 // RUN:   --bijection_outputs=reduce
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_column/vectorize_x4_s16.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_column/vectorize_x4_s16.hlo
index a708bb6f776348..33784cb7934bfb 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_column/vectorize_x4_s16.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_column/vectorize_x4_s16.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
 // RUN:   FileCheck %s
 // RUN: gpu_test_correctness %s --bijection_inputs=reduce:0 \
 // RUN:   --bijection_outputs=reduce
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_column_small/f32_2.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_column_small/f32_2.hlo
index 3d5b17ddace512..5db57916dac17b 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_column_small/f32_2.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_column_small/f32_2.hlo
@@ -1,5 +1,5 @@
 // RUN: gpu_test_correctness %s --bijection_inputs=reduce:0 --bijection_outputs=reduce
-// RUN: fusion_to_mlir %s | emitters_opt -cse -xla-simplify-arith -canonicalize | FileCheck %s --dump-input=always
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -cse -xla-simplify-arith -canonicalize | FileCheck %s --dump-input=always
 
 add {
   %p0 = f32[] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_column_small/f32_32_v2.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_column_small/f32_32_v2.hlo
index a0c2cb27189fd8..1f05c5da781720 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_column_small/f32_32_v2.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_column_small/f32_32_v2.hlo
@@ -1,5 +1,5 @@
 
-// RUN: fusion_to_mlir %s | emitters_opt -cse -xla-simplify-arith -canonicalize | FileCheck %s --dump-input=always
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -cse -xla-simplify-arith -canonicalize | FileCheck %s --dump-input=always
 
 add {
   %p0 = f32[] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_column_small/f32_8_v2.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_column_small/f32_8_v2.hlo
index 66aabaae1e3959..b0af1595363489 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_column_small/f32_8_v2.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_column_small/f32_8_v2.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -cse -xla-simplify-arith -canonicalize | FileCheck %s --dump-input=always
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -cse -xla-simplify-arith -canonicalize | FileCheck %s --dump-input=always
 // RUN: gpu_test_correctness %s --bijection_inputs=reduce:0 --bijection_outputs=reduce
 
 add {
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_column_small/s8_f32_32_v4.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_column_small/s8_f32_32_v4.hlo
index 5605fbbdb902e6..f3b92b171b7569 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_column_small/s8_f32_32_v4.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_column_small/s8_f32_32_v4.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -cse -xla-simplify-arith -canonicalize | FileCheck %s --dump-input=always
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -cse -xla-simplify-arith -canonicalize | FileCheck %s --dump-input=always
 
 add {
   %p0 = f32[] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_multirow/f16_v4.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_multirow/f16_v4.hlo
index a4317bea0e2921..7b48937d7edad4 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_multirow/f16_v4.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_multirow/f16_v4.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -xla-test-optimize \
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -xla-test-optimize \
 // RUN:   -xla-gpu-test-transform-loops | FileCheck %s
 
 // The reference implementation reduces in f64, so we need a larger tolerance.
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_multirow/f32_x8.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_multirow/f32_x8.hlo
index 90bef6674de240..140fea04100745 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_multirow/f32_x8.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_multirow/f32_x8.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | FileCheck %s
+// RUN: gpu_fusion_to_mlir %s | FileCheck %s
 // RUN: gpu_test_correctness %s --bijection_inputs=reduce:0 --bijection_outputs=reduce
 
 add {
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_multirow/f32_x8_no_inputs.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_multirow/f32_x8_no_inputs.hlo
index 82731f64a8604a..dd5db5cdc42299 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_multirow/f32_x8_no_inputs.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_multirow/f32_x8_no_inputs.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | FileCheck %s
+// RUN: gpu_fusion_to_mlir %s | FileCheck %s
 // RUN: gpu_test_correctness %s --bijection_inputs=reduce:0 --bijection_outputs=reduce
 
 add {
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_multirow/maximum_vector_size.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_multirow/maximum_vector_size.hlo
index 7a44c2b6cc1910..f8e2273cc0a0ca 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_multirow/maximum_vector_size.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_multirow/maximum_vector_size.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -xla-test-optimize \
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -xla-test-optimize \
 // RUN:   -xla-gpu-test-transform-loops | FileCheck %s
 // RUN: gpu_test_correctness %s
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_multirow/pred_mof_x2_v4.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_multirow/pred_mof_x2_v4.hlo
index be471b7c0055ff..73fc1d56d88afb 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_multirow/pred_mof_x2_v4.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_multirow/pred_mof_x2_v4.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -xla-test-optimize \
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -xla-test-optimize \
 // RUN:   -xla-gpu-test-transform-loops | FileCheck %s
 // RUN: gpu_test_correctness %s --bijection_inputs=tmp_5:0 \
 // RUN:   --bijection_inputs=tmp_8:0 --bijection_outputs=tmp_5 \
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/epilogue.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/epilogue.hlo
index 1deb8e37b7ca7d..cf0742dcc9a073 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/epilogue.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/epilogue.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -xla-test-optimize \
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -xla-test-optimize \
 // RUN:   --inline="default-pipeline='cse'" | FileCheck %s
 // RUN: gpu_test_correctness %s --bijection_inputs=reduce:0 \
 // RUN:   --bijection_outputs=reduce
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/inefficient_codegen.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/inefficient_codegen.hlo
index db2aef638f371c..cc769fad0372ec 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/inefficient_codegen.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/inefficient_codegen.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -xla-test-optimize \
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -xla-test-optimize \
 // RUN:   --inline="default-pipeline='cse'" | FileCheck %s
 // RUN: gpu_test_correctness %s --bijection_inputs=reduce1:0 \
 // RUN:   --bijection_inputs=reduce2:0 --bijection_outputs=reduce1 \
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/layout.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/layout.hlo
index 0586a28bcf6000..cc4f386d16dd87 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/layout.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/layout.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | FileCheck %s
+// RUN: gpu_fusion_to_mlir %s | FileCheck %s
 // RUN: gpu_test_correctness %s --bijection_inputs=reduce:0 --bijection_outputs=reduce
 
 add {
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/minor_and_major.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/minor_and_major.hlo
index e0b05c215bf902..b38f54822ee1e4 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/minor_and_major.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/minor_and_major.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -xla-test-optimize \
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -xla-test-optimize \
 // RUN:   --inline="default-pipeline='cse'" | FileCheck %s
 // RUN: gpu_test_correctness %s --bijection_inputs=reduce:0 \
 // RUN:   --bijection_outputs=reduce
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/mof_epilogue.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/mof_epilogue.hlo
index 78de80b882e1be..54012a6e5ed3ca 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/mof_epilogue.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/mof_epilogue.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -xla-test-optimize \
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -xla-test-optimize \
 // RUN:   --inline="default-pipeline='cse'" | FileCheck %s
 // RUN: gpu_test_correctness %s --bijection_inputs=reduce1:0 \
 // RUN:   --bijection_inputs=reduce2:0 --bijection_outputs=reduce1 \
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/reduction_groups.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/reduction_groups.hlo
index 0d17d018212808..48e0d6e6d8f0ab 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/reduction_groups.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/reduction_groups.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | FileCheck %s
+// RUN: gpu_fusion_to_mlir %s | FileCheck %s
 // RUN: gpu_test_correctness %s
 
 %add_f32 {
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/reduction_groups_two_in_same_group.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/reduction_groups_two_in_same_group.hlo
index 3cca83f157acbd..53a42478a2f58e 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/reduction_groups_two_in_same_group.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/reduction_groups_two_in_same_group.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | FileCheck %s
+// RUN: gpu_fusion_to_mlir %s | FileCheck %s
 // RUN: gpu_test_correctness %s
 
 %add_f32 {
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/side_output.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/side_output.hlo
index cebf80c93a9925..dee588f07ee3b9 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/side_output.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/side_output.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
 // RUN:   FileCheck %s
 // RUN: gpu_test_correctness %s --bijection_inputs=reduce --bijection_outputs=exp
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/side_output_s4.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/side_output_s4.hlo
index 6a9fea98096f13..daa67fc7e18fd1 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/side_output_s4.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/side_output_s4.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | FileCheck %s
+// RUN: gpu_fusion_to_mlir %s | FileCheck %s
 // RUN: gpu_test_correctness %s
 
 add {
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/small_f64.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/small_f64.hlo
index d1aa2e82926971..d84b9e1779cf1b 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/small_f64.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/small_f64.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | FileCheck %s
+// RUN: gpu_fusion_to_mlir %s | FileCheck %s
 // RUN: gpu_test_correctness %s --bijection_inputs=reduce:0 --bijection_outputs=reduce
 
 add {
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/small_s32.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/small_s32.hlo
index ee02c00172f1ad..1ffc15e673440c 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/small_s32.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/small_s32.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | FileCheck %s
+// RUN: gpu_fusion_to_mlir %s | FileCheck %s
 // RUN: gpu_test_correctness %s --bijection_inputs=reduce:0 --bijection_outputs=reduce
 
 add {
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/variadic_f32.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/variadic_f32.hlo
index da5bc2b0705b40..24f5b2258382b4 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/variadic_f32.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/reduce_row/variadic_f32.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
 // RUN:   FileCheck %s
 // RUN: gpu_test_correctness %s --bijection_inputs=reduce:0,1 \
 // RUN:   --bijection_outputs=reduce
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/add.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/add.hlo
index 113e405ed9e02e..2cd15005d85519 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/add.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/add.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
 // RUN:   FileCheck %s
 // RUN: gpu_test_correctness %s --bijection_inputs=scatter:2
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/add_vectorized.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/add_vectorized.hlo
index bfad32f2d5bb5c..e1e93d3dd68ff8 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/add_vectorized.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/add_vectorized.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -xla-test-optimize \
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -xla-test-optimize \
 // RUN:    -xla-gpu-test-transform-loops | FileCheck %s
 // RUN: gpu_test_correctness %s --bijection_inputs=scatter:2
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/overwrite.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/overwrite.hlo
index 792a8d25734fd4..877efcbc2da4b5 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/overwrite.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/overwrite.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
 // RUN:   FileCheck %s
 // RUN: gpu_test_correctness %s
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/scatter_s8.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/scatter_s8.hlo
index 7e834550393c97..ab3f23128b6c7d 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/scatter_s8.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/scatter_s8.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -xla-test-optimize --inline |\
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -xla-test-optimize --inline |\
 // RUN: FileCheck %s
 // RUN: gpu_test_correctness %s --bijection_inputs=scatter:2
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/sorted_empty_indices.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/sorted_empty_indices.hlo
index 86cb7efc8f219d..f95ef35da27a26 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/sorted_empty_indices.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/sorted_empty_indices.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s   | emitters_opt -xla-test-optimize \
+// RUN: gpu_fusion_to_mlir %s   | emitters_opt -xla-test-optimize \
 // RUN:   -xla-gpu-test-transform-loops  | FileCheck %s
 // RUN: gpu_test_correctness %s --bijection_inputs=scatter:2
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/sorted_indices.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/sorted_indices.hlo
index 6bd0c8261ed5eb..3a79e83013b49c 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/sorted_indices.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/sorted_indices.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s   | emitters_opt -xla-test-optimize \
+// RUN: gpu_fusion_to_mlir %s   | emitters_opt -xla-test-optimize \
 // RUN:   -xla-gpu-test-transform-loops  | FileCheck %s
 // RUN: gpu_test_correctness %s --bijection_inputs=scatter:2
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/sorted_indices_large.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/sorted_indices_large.hlo
index c110192edcc378..5ef3af7e8f0115 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/sorted_indices_large.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/sorted_indices_large.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s   | emitters_opt -xla-test-optimize \
+// RUN: gpu_fusion_to_mlir %s   | emitters_opt -xla-test-optimize \
 // RUN:   -xla-gpu-test-transform-loops  | FileCheck %s
 
 mul {
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/sorted_indices_small.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/sorted_indices_small.hlo
index ad85709db9cf89..f9babbcd4bdd12 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/sorted_indices_small.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/sorted_indices_small.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s   | emitters_opt -xla-test-optimize \
+// RUN: gpu_fusion_to_mlir %s   | emitters_opt -xla-test-optimize \
 // RUN: | FileCheck %s
 // RUN: gpu_test_correctness %s --bijection_inputs=scatter:2
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/sorted_unique_indices.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/sorted_unique_indices.hlo
index c226ea4955a294..bad9ddd3edb2f3 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/sorted_unique_indices.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/sorted_unique_indices.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s   | emitters_opt -xla-test-optimize \
+// RUN: gpu_fusion_to_mlir %s   | emitters_opt -xla-test-optimize \
 // RUN:   -xla-gpu-test-transform-loops  | FileCheck %s
 // RUN: gpu_test_correctness %s --bijection_inputs=scatter:2
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/unique_indices.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/unique_indices.hlo
index 2004c4b94a02af..9294802dafd008 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/unique_indices.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/unique_indices.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
 // RUN:   FileCheck %s
 // RUN: gpu_test_correctness %s
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/unsigned.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/unsigned.hlo
index 65dcc1b7472d44..d19e1815c9ca67 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/unsigned.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/unsigned.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
 // RUN:   FileCheck %s
 // RUN: gpu_test_correctness %s
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/variadic_scatter_argmax.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/variadic_scatter_argmax.hlo
index f9e3bbcf532726..7ff52e4a00999a 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/variadic_scatter_argmax.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/variadic_scatter_argmax.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -xla-test-optimize | FileCheck %s
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -xla-test-optimize | FileCheck %s
 // RUN: gpu_test_correctness %s --bijection_inputs=scatter:3
 
 argmax_combiner {
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/variadic_scatter_simple.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/variadic_scatter_simple.hlo
index e1dd063a84a730..e93c31b726ef08 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/variadic_scatter_simple.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/scatter/variadic_scatter_simple.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -xla-test-optimize | FileCheck %s
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -xla-test-optimize | FileCheck %s
 // RUN: gpu_test_correctness %s --bijection_inputs=scatter:3
 
 override {
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/epilogue.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/epilogue.hlo
index 88faae011952b0..aeaa6681d64c8c 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/epilogue.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/epilogue.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
 // RUN:   FileCheck %s
 // RUN: gpu_test_correctness %s
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/fused_transpose_021.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/fused_transpose_021.hlo
index dfa915a1a240ea..aa747beedf475b 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/fused_transpose_021.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/fused_transpose_021.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
 // RUN:   FileCheck %s
 // RUN: gpu_test_correctness %s
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/fused_transpose_102.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/fused_transpose_102.hlo
index 6cb6ab5f93c482..dc939813373495 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/fused_transpose_102.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/fused_transpose_102.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
 // RUN:   FileCheck %s
 // RUN: gpu_test_correctness %s
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/fused_transpose_210.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/fused_transpose_210.hlo
index 1488e28bb9f310..78d83b446d84a8 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/fused_transpose_210.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/fused_transpose_210.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
 // RUN:   FileCheck %s
 // RUN: gpu_test_correctness %s
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/mixed_indexing.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/mixed_indexing.hlo
index 54e9be778cc392..79b43f020d165f 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/mixed_indexing.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/mixed_indexing.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
 // RUN:   FileCheck %s
 // RUN: gpu_test_correctness %s
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/multiple_roots.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/multiple_roots.hlo
index d57bc16386c9e8..ff1797ac36b003 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/multiple_roots.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/multiple_roots.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
 // RUN:   FileCheck %s
 // RUN: gpu_test_correctness %s
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/multiple_roots_mixed_rank.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/multiple_roots_mixed_rank.hlo
index 11114c8d481484..32e603e2e31ee6 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/multiple_roots_mixed_rank.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/multiple_roots_mixed_rank.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | FileCheck %s
+// RUN: gpu_fusion_to_mlir %s | FileCheck %s
 // RUN: gpu_test_correctness %s
 
 HloModule module
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/multiple_roots_one_shmem_transpose.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/multiple_roots_one_shmem_transpose.hlo
index 90846fcc7a7541..3b30f72dd5d382 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/multiple_roots_one_shmem_transpose.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/multiple_roots_one_shmem_transpose.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
 // RUN:   FileCheck %s
 // RUN: gpu_test_correctness %s
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/packed_transpose_bf16.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/packed_transpose_bf16.hlo
index b74462c0f40b2c..2a25efa4ad5215 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/packed_transpose_bf16.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/packed_transpose_bf16.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
 // RUN:   FileCheck %s
 // RUN: gpu_test_correctness %s
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/packed_transpose_f16.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/packed_transpose_f16.hlo
index 53f083a500c383..44fd1c3d07721d 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/packed_transpose_f16.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/packed_transpose_f16.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
 // RUN:   FileCheck %s
 // RUN: gpu_test_correctness %s
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/packed_transpose_multiple_heroes.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/packed_transpose_multiple_heroes.hlo
index 170a4c80940d67..91f45ef546d211 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/packed_transpose_multiple_heroes.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/packed_transpose_multiple_heroes.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
 // RUN:   FileCheck %s --dump-input=fail
 // RUN: gpu_test_correctness %s
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/packed_transpose_multiple_roots.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/packed_transpose_multiple_roots.hlo
index a760db83238063..639b03399bd9dd 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/packed_transpose_multiple_roots.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/packed_transpose_multiple_roots.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
 // RUN:   FileCheck %s
 // RUN: gpu_test_correctness %s
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/packed_transpose_s16.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/packed_transpose_s16.hlo
index 500a98c3b2ce79..90e3b7281007d9 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/packed_transpose_s16.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/packed_transpose_s16.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -xla-test-optimize --inline |\
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -xla-test-optimize --inline |\
 // RUN:   FileCheck %s
 // RUN: gpu_test_correctness %s
 fusion {
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/packed_transpose_s4.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/packed_transpose_s4.hlo
index 88dd4caa54b968..cc35fe62474ac8 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/packed_transpose_s4.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/packed_transpose_s4.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
 // RUN:   FileCheck %s
 // RUN: gpu_test_correctness %s
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/packed_transpose_s8.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/packed_transpose_s8.hlo
index f8f309951dfe21..2fdfb58165bca8 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/packed_transpose_s8.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/packed_transpose_s8.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
 // RUN:   FileCheck %s
 // RUN: gpu_test_correctness %s
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/packed_transpose_side_output.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/packed_transpose_side_output.hlo
index c82d53d2d52e63..73749aa90678af 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/packed_transpose_side_output.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/packed_transpose_side_output.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
 // RUN:   FileCheck %s
 // RUN: gpu_test_correctness %s
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/packed_transpose_two_heroes.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/packed_transpose_two_heroes.hlo
index fbf32f148feff1..09ed1c0ce64df5 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/packed_transpose_two_heroes.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/packed_transpose_two_heroes.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
 // RUN:   FileCheck %s
 // RUN: gpu_test_correctness %s
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/partial_tile.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/partial_tile.hlo
index 86c29b45f38a61..9e18e084830bbe 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/partial_tile.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/partial_tile.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
 // RUN:   FileCheck %s
 // RUN: gpu_test_correctness %s
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/partial_tile_s4.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/partial_tile_s4.hlo
index 5240f69cbc070e..28322bcbd2b3ed 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/partial_tile_s4.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/partial_tile_s4.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
 // RUN:   FileCheck %s
 // RUN: gpu_test_correctness %s
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/side_outputs.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/side_outputs.hlo
index 707862fee649c6..008c97daf3009c 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/side_outputs.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/side_outputs.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
 // RUN:   FileCheck %s
 // RUN: gpu_test_correctness %s
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/side_outputs_inplace.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/side_outputs_inplace.hlo
index 1247d49681155c..2f3b572ba284f5 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/side_outputs_inplace.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/side_outputs_inplace.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
 // RUN:   FileCheck %s
 // RUN: gpu_test_correctness %s
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/transpose_10.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/transpose_10.hlo
index 8496dcab4c62f1..824e16ea34faa1 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/transpose_10.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/transpose_10.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
 // RUN:   FileCheck %s
 // RUN: gpu_test_correctness %s --bijection_inputs=transpose --bijection_outputs=transpose
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/transpose_102_s4.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/transpose_102_s4.hlo
index 7ff8a110d38ce0..6fc72e80595e96 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/transpose_102_s4.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/transpose_102_s4.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
 // RUN:   FileCheck %s
 // RUN: gpu_test_correctness %s
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/transpose_10_s4.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/transpose_10_s4.hlo
index 666191de1dd7a3..ca2463c2fa3fd8 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/transpose_10_s4.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/transpose_10_s4.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
 // RUN:   FileCheck %s
 // RUN: gpu_test_correctness %s --bijection_inputs=transpose --bijection_outputs=transpose
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/transpose_1302.hlo b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/transpose_1302.hlo
index ada87225690890..6f0abcd6a1561b 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/transpose_1302.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/transpose/transpose_1302.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
+// RUN: gpu_fusion_to_mlir %s | emitters_opt -xla-test-optimize |\
 // RUN:   FileCheck %s
 // RUN: gpu_test_correctness %s
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/kernels/BUILD b/third_party/xla/xla/backends/gpu/codegen/kernels/BUILD
index 377cfd6d98f2a5..586cbc804bbe00 100644
--- a/third_party/xla/xla/backends/gpu/codegen/kernels/BUILD
+++ b/third_party/xla/xla/backends/gpu/codegen/kernels/BUILD
@@ -90,8 +90,6 @@ cc_library(
     visibility = [":friends"],
     deps = [
         ":custom_kernel",
-        "//xla/stream_executor:kernel",
-        "//xla/stream_executor:kernel_args",
         "//xla/stream_executor:kernel_args_packing_spec",
         "//xla/stream_executor:kernel_spec",
         "//xla/stream_executor:launch_dim",
diff --git a/third_party/xla/xla/backends/gpu/codegen/kernels/ptx_custom_kernel.cc b/third_party/xla/xla/backends/gpu/codegen/kernels/ptx_custom_kernel.cc
index 44883901699c6d..3f0a0b716a6105 100644
--- a/third_party/xla/xla/backends/gpu/codegen/kernels/ptx_custom_kernel.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/kernels/ptx_custom_kernel.cc
@@ -17,6 +17,7 @@ limitations under the License.
 
 #include <cstddef>
 #include <cstdint>
+#include <memory>
 #include <string>
 #include <utility>
 #include <vector>
@@ -24,8 +25,6 @@ limitations under the License.
 #include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
 #include "xla/backends/gpu/codegen/kernels/custom_kernel.h"
-#include "xla/stream_executor/kernel.h"
-#include "xla/stream_executor/kernel_args.h"
 #include "xla/stream_executor/kernel_args_packing_spec.h"
 #include "xla/stream_executor/kernel_spec.h"
 #include "xla/stream_executor/launch_dim.h"
@@ -86,4 +85,16 @@ absl::StatusOr<CustomKernel> CreateOwnedCubinCustomKernel(
                       thread_dim, shared_memory_bytes);
 }
 
+absl::StatusOr<CustomKernel> CreateSharedCubinCustomKernel(
+    std::string kernel_name, std::shared_ptr<const std::vector<uint8_t>> cubin,
+    int num_args, se::BlockDim block_dim, se::ThreadDim thread_dim,
+    size_t shared_memory_bytes) {
+  se::KernelLoaderSpec kernel_spec =
+      se::KernelLoaderSpec::CreateSharedCudaCubinInMemorySpec(
+          std::move(cubin), kernel_name, /*arity=*/num_args,
+          se::KernelArgsPackingSpec::Identity(num_args));
+  return CustomKernel(std::move(kernel_name), std::move(kernel_spec), block_dim,
+                      thread_dim, shared_memory_bytes);
+}
+
 }  // namespace xla::gpu::kernel
diff --git a/third_party/xla/xla/backends/gpu/codegen/kernels/ptx_custom_kernel.h b/third_party/xla/xla/backends/gpu/codegen/kernels/ptx_custom_kernel.h
index 8ecc0a7cc1fcf0..4fef15cc4c3ea6 100644
--- a/third_party/xla/xla/backends/gpu/codegen/kernels/ptx_custom_kernel.h
+++ b/third_party/xla/xla/backends/gpu/codegen/kernels/ptx_custom_kernel.h
@@ -17,7 +17,10 @@ limitations under the License.
 #define XLA_BACKENDS_GPU_CODEGEN_KERNELS_PTX_CUSTOM_KERNEL_H_
 
 #include <cstddef>
+#include <cstdint>
+#include <memory>
 #include <string>
+#include <vector>
 
 #include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
@@ -48,5 +51,16 @@ absl::StatusOr<CustomKernel> CreateOwnedCubinCustomKernel(
     se::BlockDim block_dim, se::ThreadDim thread_dim,
     size_t shared_memory_bytes);
 
+// Like CreateOwnedCubinCustomKernel, but the CUBIN buffer is shared instead of
+// copied. Prefer this whenever the CUBIN is already held in a reference counted
+// buffer (e.g. it comes from the KernelReuseCache), so that a kernel that is
+// invoked multiple times in a module is only stored once.
+//
+// `cubin` must not be null.
+absl::StatusOr<CustomKernel> CreateSharedCubinCustomKernel(
+    std::string kernel_name, std::shared_ptr<const std::vector<uint8_t>> cubin,
+    int num_args, se::BlockDim block_dim, se::ThreadDim thread_dim,
+    size_t shared_memory_bytes);
+
 }  // namespace xla::gpu::kernel
 #endif  // XLA_BACKENDS_GPU_CODEGEN_KERNELS_PTX_CUSTOM_KERNEL_H_
diff --git a/third_party/xla/xla/backends/gpu/codegen/tools/BUILD b/third_party/xla/xla/backends/gpu/codegen/tools/BUILD
index 6ad8aafadc13ed..b14db5fafbf8f4 100644
--- a/third_party/xla/xla/backends/gpu/codegen/tools/BUILD
+++ b/third_party/xla/xla/backends/gpu/codegen/tools/BUILD
@@ -65,7 +65,7 @@ cc_library(
 )
 
 xla_cc_binary(
-    name = "fusion_to_mlir",
+    name = "gpu_fusion_to_mlir",
     testonly = 1,
     srcs = ["fusion_to_mlir.cc"],
     # We want to use this tool for lit tests. Due to hermetic cuda, we need to
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/fusion.cc b/third_party/xla/xla/backends/gpu/codegen/triton/fusion.cc
index f156c165762938..75d27fc3e4fba4 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/fusion.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/fusion.cc
@@ -108,9 +108,12 @@ absl::StatusOr<EmitArgs> EmitCollectiveFusion(
     ABSL_ASSIGN_OR_RETURN(CollectiveKernelSpec kernel_spec,
                      CreateCollectiveKernelSpec(
                          fusion_instr, result.entry.launch_dimensions));
-    auto cubin = result.entry.binary.empty()
-                     ? std::nullopt
-                     : std::make_optional(std::move(result.entry.binary));
+    // `CollectiveKernelThunk` owns its cubin, so we have to copy it out of the
+    // shared buffer here.
+    std::optional<std::vector<uint8_t>> cubin;
+    if (result.entry.binary != nullptr && !result.entry.binary->empty()) {
+      cubin = *result.entry.binary;
+    }
     return ThunkSequence::Of<CollectiveKernelThunk>(
         std::move(info), config, std::move(kernel_spec), std::move(buffers),
         /*is_collective_kernel_enabled=*/true, result.entry.kernel_name,
@@ -195,8 +198,9 @@ AsyncThunkSequence TritonFusion::Emit(
             EmitResult result) -> absl::StatusOr<ThunkSequence> {
       ABSL_ASSIGN_OR_RETURN(
           CustomKernel custom_kernel,
-          kernel::CreateOwnedCubinCustomKernel(
-              result.entry.kernel_name, result.entry.binary,
+          kernel::CreateSharedCubinCustomKernel(
+              std::move(result.entry.kernel_name),
+              std::move(result.entry.binary),
               result.kernel_arguments.args().size(),
               result.entry.launch_dimensions.block_counts(),
               result.entry.launch_dimensions.thread_counts_per_block(),
@@ -204,7 +208,7 @@ AsyncThunkSequence TritonFusion::Emit(
       return ThunkSequence::Of<CustomKernelThunk>(
           thunk_info, std::move(custom_kernel), result.kernel_arguments,
           result.entry.use_pdl, std::vector<int64_t>{},
-          result.entry.tma_metadata);
+          std::move(result.entry.tma_metadata));
     };
     emit_args = {
         std::move(make_thunk),
@@ -320,13 +324,14 @@ xla::Future<TritonFusion::EmitResult> TritonFusion::Emit(
                     shmem_bytes = triton_wrapper_result.shmem_bytes,
                     use_pdl = triton_wrapper_result.use_pdl](
                        const std::vector<uint8_t>& cubin) mutable {
-                return KernelReuseCache::Entry{std::move(kernel_name),
-                                               launch_dims,
-                                               /*cluster_dim=*/std::nullopt,
-                                               shmem_bytes,
-                                               cubin,
-                                               tma_metadata,
-                                               use_pdl};
+                return KernelReuseCache::Entry{
+                    std::move(kernel_name),
+                    launch_dims,
+                    /*cluster_dim=*/std::nullopt,
+                    shmem_bytes,
+                    std::make_shared<const std::vector<uint8_t>>(cubin),
+                    tma_metadata,
+                    use_pdl};
               });
         });
   };
@@ -336,9 +341,9 @@ xla::Future<TritonFusion::EmitResult> TritonFusion::Emit(
           hlo_computation, kernel_arguments.args(),
           /*discriminator=*/"TritonFusion", generate);
   return status_or_entry.Map([kernel_arguments = std::move(kernel_arguments)](
-                                 const KernelReuseCache::Entry* entry) mutable
+                                 const KernelReuseCache::Entry& entry) mutable
                                  -> absl::StatusOr<EmitResult> {
-    return EmitResult{*entry, std::move(kernel_arguments)};
+    return EmitResult{entry, std::move(kernel_arguments)};
   });
 }
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/BUILD b/third_party/xla/xla/backends/gpu/codegen/triton/tests/BUILD
index e241fc166c2552..7459d9e0139743 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/BUILD
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/BUILD
@@ -143,9 +143,10 @@ xla_cc_test(
         "//xla/stream_executor:device_description",
         "//xla/stream_executor/cuda:cuda_compute_capability",
         "//xla/tests:xla_internal_test_main",
-        "//xla/tsl/platform:statusor",
         "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_macros",
         "@com_google_absl//absl/status:status_matchers",
+        "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:string_view",
         "@com_google_googletest//:gtest",
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_device_test.cc b/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_device_test.cc
index e0e1be21c78b8b..dbf9f96b027714 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_device_test.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_device_test.cc
@@ -3606,6 +3606,277 @@ INSTANTIATE_TEST_SUITE_P(
                        ::testing::Bool(), ::testing::Bool(), ::testing::Bool()),
     ScaledDotCoverageTestParamToString);
 
+// Scans are only emitted with the experimental tiling propagation enabled.
+class ScanTritonEmitterTest : public TritonEmitterTest {
+ public:
+  bool EnableTilingPropagation() const override { return true; }
+
+ protected:
+  // Returns arguments for the entry computation of `module` filled with small
+  // integers. All intermediate scan results are then exactly representable in
+  // f32, so that the results match the interpreter reference bit-exactly,
+  // independently of the order in which the scan combines the elements.
+  absl::StatusOr<std::vector<Literal>> MakeSmallIntegerArguments(
+      const HloModule& module) {
+    std::vector<Literal> arguments;
+    for (const HloInstruction* parameter :
+         module.entry_computation()->parameter_instructions()) {
+      Literal literal(parameter->shape());
+      ABSL_RETURN_IF_ERROR(
+          literal.Populate<float>([](absl::Span<const int64_t> indices) {
+            int64_t value = 0;
+            for (int64_t index : indices) {
+              value += index;
+            }
+            return static_cast<float>(value % 7 - 3);
+          }));
+      arguments.push_back(std::move(literal));
+    }
+    return arguments;
+  }
+
+  void RunAndCompareScan(absl::string_view hlo_text) {
+    ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
+                         ParseAndReturnVerifiedModule(hlo_text));
+    ASSERT_OK_AND_ASSIGN(std::vector<Literal> arguments,
+                         MakeSmallIntegerArguments(*module));
+    EXPECT_TRUE(RunAndCompareNoHloPasses(
+        std::move(module), LiteralUtil::MakePointers(arguments), kExactMatch));
+  }
+};
+
+TEST_F(ScanTritonEmitterTest, SingleTileScan) {
+  constexpr absl::string_view kHloText = R"(
+add_computation {
+  p0 = f32[] parameter(0)
+  p1 = f32[] parameter(1)
+  add = f32[] add(p0, p1)
+  ROOT tuple = (f32[], f32[]) tuple(add, add)
+}
+
+scan_fusion {
+  p0 = f32[1024] parameter(0)
+  p1 = f32[] parameter(1)
+  scan = (f32[1024], f32[]) scan(p0, p1), dimensions={0}, num_carries=1,
+    is_associative=true, to_apply=add_computation
+  ROOT get-tuple-element = f32[1024] get-tuple-element(scan), index=0
+}
+
+ENTRY entry_computation {
+  p0 = f32[1024] parameter(0)
+  p1 = f32[] parameter(1)
+  ROOT fusion = f32[1024] fusion(p0, p1), kind=kCustom, calls=scan_fusion,
+    backend_config={"fusion_backend_config":{"kind":"__triton",
+      "block_level_fusion_config":{"output_tiles":[{"sizes":[]}],
+        "num_warps":4,"num_ctas":1,"num_stages":1}}}
+})";
+  RunAndCompareScan(kHloText);
+}
+
+TEST_F(ScanTritonEmitterTest, TiledScan) {
+  constexpr absl::string_view kHloText = R"(
+add_computation {
+  p0 = f32[] parameter(0)
+  p1 = f32[] parameter(1)
+  add = f32[] add(p0, p1)
+  ROOT tuple = (f32[], f32[]) tuple(add, add)
+}
+
+scan_fusion {
+  p0 = f32[1024] parameter(0)
+  p1 = f32[] parameter(1)
+  scan = (f32[1024], f32[]) scan(p0, p1), dimensions={0}, num_carries=1,
+    is_associative=true, to_apply=add_computation, backend_config={sizes:[256]}
+  ROOT get-tuple-element = f32[1024] get-tuple-element(scan), index=0
+}
+
+ENTRY entry_computation {
+  p0 = f32[1024] parameter(0)
+  p1 = f32[] parameter(1)
+  ROOT fusion = f32[1024] fusion(p0, p1), kind=kCustom, calls=scan_fusion,
+    backend_config={"fusion_backend_config":{"kind":"__triton",
+      "block_level_fusion_config":{"output_tiles":[{"sizes":[]}],
+        "num_warps":4,"num_ctas":1,"num_stages":1}}}
+})";
+  RunAndCompareScan(kHloText);
+}
+
+// The scan dimension is not divisible by the tile size, so the input tile of
+// the last loop iteration extends past the dimension bound and must be masked
+// before it is fed to the scan.
+TEST_F(ScanTritonEmitterTest, TiledScanWithPartialTile) {
+  constexpr absl::string_view kHloText = R"(
+add_computation {
+  p0 = f32[] parameter(0)
+  p1 = f32[] parameter(1)
+  add = f32[] add(p0, p1)
+  ROOT tuple = (f32[], f32[]) tuple(add, add)
+}
+
+scan_fusion {
+  p0 = f32[1000] parameter(0)
+  p1 = f32[] parameter(1)
+  scan = (f32[1000], f32[]) scan(p0, p1), dimensions={0}, num_carries=1,
+    is_associative=true, to_apply=add_computation, backend_config={sizes:[256]}
+  ROOT get-tuple-element = f32[1000] get-tuple-element(scan), index=0
+}
+
+ENTRY entry_computation {
+  p0 = f32[1000] parameter(0)
+  p1 = f32[] parameter(1)
+  ROOT fusion = f32[1000] fusion(p0, p1), kind=kCustom, calls=scan_fusion,
+    backend_config={"fusion_backend_config":{"kind":"__triton",
+      "block_level_fusion_config":{"output_tiles":[{"sizes":[]}],
+        "num_warps":4,"num_ctas":1,"num_stages":1}}}
+})";
+  RunAndCompareScan(kHloText);
+}
+
+TEST_F(ScanTritonEmitterTest, TiledScanOfMinorDimension) {
+  constexpr absl::string_view kHloText = R"(
+add_computation {
+  p0 = f32[16] parameter(0)
+  p1 = f32[16] parameter(1)
+  add = f32[16] add(p0, p1)
+  ROOT tuple = (f32[16], f32[16]) tuple(add, add)
+}
+
+scan_fusion {
+  p0 = f32[16,1024] parameter(0)
+  p1 = f32[16] parameter(1)
+  scan = (f32[16,1024], f32[16]) scan(p0, p1), dimensions={1}, num_carries=1,
+    is_associative=true, to_apply=add_computation, backend_config={sizes:[256]}
+  ROOT get-tuple-element = f32[16,1024] get-tuple-element(scan), index=0
+}
+
+ENTRY entry_computation {
+  p0 = f32[16,1024] parameter(0)
+  p1 = f32[16] parameter(1)
+  ROOT fusion = f32[16,1024] fusion(p0, p1), kind=kCustom, calls=scan_fusion,
+    backend_config={"fusion_backend_config":{"kind":"__triton",
+      "block_level_fusion_config":{"output_tiles":[{"sizes":[16]}],
+        "num_warps":4,"num_ctas":1,"num_stages":1}}}
+})";
+  RunAndCompareScan(kHloText);
+}
+
+// Scanning the major dimension distributes the scan dimension across warps
+// differently than the minor one, which changes how the carry reduction is
+// lowered.
+TEST_F(ScanTritonEmitterTest, TiledScanOfMajorDimension) {
+  constexpr absl::string_view kHloText = R"(
+add_computation {
+  p0 = f32[16] parameter(0)
+  p1 = f32[16] parameter(1)
+  add = f32[16] add(p0, p1)
+  ROOT tuple = (f32[16], f32[16]) tuple(add, add)
+}
+
+scan_fusion {
+  p0 = f32[1024,16] parameter(0)
+  p1 = f32[16] parameter(1)
+  scan = (f32[1024,16], f32[16]) scan(p0, p1), dimensions={0}, num_carries=1,
+    is_associative=true, to_apply=add_computation, backend_config={sizes:[256]}
+  ROOT get-tuple-element = f32[1024,16] get-tuple-element(scan), index=0
+}
+
+ENTRY entry_computation {
+  p0 = f32[1024,16] parameter(0)
+  p1 = f32[16] parameter(1)
+  ROOT fusion = f32[1024,16] fusion(p0, p1), kind=kCustom, calls=scan_fusion,
+    backend_config={"fusion_backend_config":{"kind":"__triton",
+      "block_level_fusion_config":{"output_tiles":[{"sizes":[16]}],
+        "num_warps":4,"num_ctas":1,"num_stages":1}}}
+})";
+  RunAndCompareScan(kHloText);
+}
+
+TEST_F(ScanTritonEmitterTest, SingleTileReverseScan) {
+  constexpr absl::string_view kHloText = R"(
+add_computation {
+  p0 = f32[] parameter(0)
+  p1 = f32[] parameter(1)
+  add = f32[] add(p0, p1)
+  ROOT tuple = (f32[], f32[]) tuple(add, add)
+}
+
+scan_fusion {
+  p0 = f32[1024] parameter(0)
+  p1 = f32[] parameter(1)
+  scan = (f32[1024], f32[]) scan(p0, p1), dimensions={0}, is_reverse=true,
+    num_carries=1, is_associative=true, to_apply=add_computation
+  ROOT get-tuple-element = f32[1024] get-tuple-element(scan), index=0
+}
+
+ENTRY entry_computation {
+  p0 = f32[1024] parameter(0)
+  p1 = f32[] parameter(1)
+  ROOT fusion = f32[1024] fusion(p0, p1), kind=kCustom, calls=scan_fusion,
+    backend_config={"fusion_backend_config":{"kind":"__triton",
+      "block_level_fusion_config":{"output_tiles":[{"sizes":[]}],
+        "num_warps":4,"num_ctas":1,"num_stages":1}}}
+})";
+  RunAndCompareScan(kHloText);
+}
+
+TEST_F(ScanTritonEmitterTest, TiledReverseScan) {
+  constexpr absl::string_view kHloText = R"(
+add_computation {
+  p0 = f32[] parameter(0)
+  p1 = f32[] parameter(1)
+  add = f32[] add(p0, p1)
+  ROOT tuple = (f32[], f32[]) tuple(add, add)
+}
+
+scan_fusion {
+  p0 = f32[1024] parameter(0)
+  p1 = f32[] parameter(1)
+  scan = (f32[1024], f32[]) scan(p0, p1), dimensions={0}, is_reverse=true,
+    num_carries=1, is_associative=true, to_apply=add_computation,
+    backend_config={sizes:[256]}
+  ROOT get-tuple-element = f32[1024] get-tuple-element(scan), index=0
+}
+
+ENTRY entry_computation {
+  p0 = f32[1024] parameter(0)
+  p1 = f32[] parameter(1)
+  ROOT fusion = f32[1024] fusion(p0, p1), kind=kCustom, calls=scan_fusion,
+    backend_config={"fusion_backend_config":{"kind":"__triton",
+      "block_level_fusion_config":{"output_tiles":[{"sizes":[]}],
+        "num_warps":4,"num_ctas":1,"num_stages":1}}}
+})";
+  RunAndCompareScan(kHloText);
+}
+
+TEST_F(ScanTritonEmitterTest, TiledReverseScanWithPartialTile) {
+  constexpr absl::string_view kHloText = R"(
+add_computation {
+  p0 = f32[] parameter(0)
+  p1 = f32[] parameter(1)
+  add = f32[] add(p0, p1)
+  ROOT tuple = (f32[], f32[]) tuple(add, add)
+}
+
+scan_fusion {
+  p0 = f32[1000] parameter(0)
+  p1 = f32[] parameter(1)
+  scan = (f32[1000], f32[]) scan(p0, p1), dimensions={0}, is_reverse=true,
+    num_carries=1, is_associative=true, to_apply=add_computation,
+    backend_config={sizes:[256]}
+  ROOT get-tuple-element = f32[1000] get-tuple-element(scan), index=0
+}
+
+ENTRY entry_computation {
+  p0 = f32[1000] parameter(0)
+  p1 = f32[] parameter(1)
+  ROOT fusion = f32[1000] fusion(p0, p1), kind=kCustom, calls=scan_fusion,
+    backend_config={"fusion_backend_config":{"kind":"__triton",
+      "block_level_fusion_config":{"output_tiles":[{"sizes":[]}],
+        "num_warps":4,"num_ctas":1,"num_stages":1}}}
+})";
+  RunAndCompareScan(kHloText);
+}
+
 }  // namespace
 }  // namespace gpu
 }  // namespace xla
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_deviceless_test.cc b/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_deviceless_test.cc
index c7937b9ceb9ce8..7244a8c7eecb78 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_deviceless_test.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_deviceless_test.cc
@@ -19,7 +19,9 @@ limitations under the License.
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include "absl/status/status.h"
+#include "absl/status/status_macros.h"
 #include "absl/status/status_matchers.h"
+#include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
 #include "absl/strings/substitute.h"
 #include "llvm/IR/LLVMContext.h"
@@ -40,7 +42,6 @@ limitations under the License.
 #include "xla/service/gpu/target_constants.h"
 #include "xla/stream_executor/cuda/cuda_compute_capability.h"
 #include "xla/stream_executor/device_description.h"
-#include "xla/tsl/platform/statusor.h"
 #include "xla/xla.pb.h"
 #include "xla/xla_data.pb.h"
 
@@ -69,6 +70,40 @@ class ExperimentalTilingTritonEmitterTest : public TritonEmitterDevicelessTest {
     debug_options.set_xla_gpu_experimental_enable_tiling_propagation(true);
     return debug_options;
   }
+
+ protected:
+  // Emits and compiles the Triton fusion that is the root of `hlo_text` and
+  // returns the Triton IR before the conversion to the TritonGPU dialect.
+  absl::StatusOr<std::string> EmitAndCompileTritonFusion(
+      absl::string_view hlo_text) {
+    ABSL_ASSIGN_OR_RETURN(std::unique_ptr<VerifiedHloModule> hlo_module,
+                     ParseAndReturnVerifiedModule(hlo_text));
+    const HloFusionInstruction& fusion = *Cast<HloFusionInstruction>(
+        hlo_module->entry_computation()->root_instruction());
+    ABSL_ASSIGN_OR_RETURN(GpuBackendConfig backend_config,
+                     fusion.backend_config<GpuBackendConfig>());
+    BlockLevelParameters block_level_parameters =
+        BlockLevelParameters::FromBlockLevelFusionConfig(
+            backend_config.fusion_backend_config().block_level_fusion_config());
+    const se::DeviceDescription dev_info =
+        TestGpuDeviceInfo::RTXA6000DeviceInfo();
+    mlir::MLIRContext mlir_context;
+    RegisterSymbolicExprStorage(&mlir_context);
+
+    ABSL_ASSIGN_OR_RETURN(TritonKernelSource triton_source,
+                     CreateTritonModule("test_fn", fusion, dev_info,
+                                        block_level_parameters, mlir_context));
+    absl::Status status =
+        TritonWrapper("test_fn", fusion, dev_info.gpu_compute_capability(),
+                      dev_info, block_level_parameters,
+                      llvm::Triple(nvptx::TargetTriple()), nvptx::DataLayout(),
+                      mlir_context)
+            .status();
+    if (!status.ok()) {
+      return status;
+    }
+    return triton_source.ToString();
+  }
 };
 
 TEST_F(TritonEmitterDevicelessTest, FailsGracefullyIfNumWarpsIsMissing) {
@@ -268,7 +303,7 @@ ENTRY entry {
 }
 
 TEST_F(ExperimentalTilingTritonEmitterTest, ScanEmitOk) {
-  const std::string kHloText = R"(
+  constexpr absl::string_view kHloText = R"(
 scan_computation {
   p_carry = f32[] parameter(0)
   p_input = f32[] parameter(1)
@@ -300,44 +335,260 @@ ENTRY main {
 }
 )";
 
-  ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> hlo_module,
-                       ParseAndReturnVerifiedModule(kHloText));
-  const HloFusionInstruction* triton_fusion = Cast<HloFusionInstruction>(
-      hlo_module->entry_computation()->root_instruction());
-  const se::DeviceDescription dev_info =
-      TestGpuDeviceInfo::RTXA6000DeviceInfo();
-  mlir::MLIRContext mlir_context;
-  RegisterSymbolicExprStorage(&mlir_context);
-  ASSERT_OK_AND_ASSIGN(auto backend_config,
-                       triton_fusion->backend_config<GpuBackendConfig>());
-
-  ASSERT_OK_AND_ASSIGN(
-      TritonKernelSource triton_source,
-      CreateTritonModule("test_fn", *triton_fusion, dev_info,
-                         BlockLevelParameters::FromBlockLevelFusionConfig(
-                             backend_config.fusion_backend_config()
-                                 .block_level_fusion_config()),
-                         mlir_context));
-
-  std::string triton_mlir = triton_source.ToString();
+  ASSERT_OK_AND_ASSIGN(std::string triton_ir,
+                       EmitAndCompileTritonFusion(kHloText));
 
+  // The scan fits into a single tile, so there is no loop and no carry. The
+  // unit-dimension init is broadcast to the tile shape, which folds into the
+  // splat of the scalar init value.
   constexpr absl::string_view kPattern = R"(
 // CHECK-LABEL: @test_fn
 // CHECK:         %[[INPUT:.*]] = xtile.extract %arg0[%c0] [1024] [1] : memref<1024xf32> -> tensor<1024xf32>
 // CHECK:         %[[INIT:.*]] = xtile.extract %arg1[] [] [] : memref<f32> -> tensor<f32>
+// CHECK:         %[[INIT_SCALAR:.*]] = tensor.extract %[[INIT]][] : tensor<f32>
 // CHECK:         %[[SCAN:.*]] = "tt.scan"(%[[INPUT]]) <{axis = 0 : i32, reverse = false}> ({
 // CHECK:         ^bb0(%[[LHS:.*]]: f32, %[[RHS:.*]]: f32):
 // CHECK:           %[[ADD:.*]] = arith.addf %[[LHS]], %[[RHS]] : f32
 // CHECK:           tt.scan.return %[[ADD]] : f32
 // CHECK:         }) : (tensor<1024xf32>) -> tensor<1024xf32>
-// CHECK:         %[[BCAST_INIT:.*]] = stablehlo.broadcast_in_dim %[[INIT]], dims = [] : (tensor<f32>) -> tensor<1024xf32>
+// CHECK:         %[[BCAST_INIT:.*]] = tt.splat %[[INIT_SCALAR]] : f32 -> tensor<1024xf32>
 // CHECK:         %[[OUTPUT:.*]] = arith.addf %[[BCAST_INIT]], %[[SCAN]] : tensor<1024xf32>
+// CHECK-NOT:     tt.reshape
 // CHECK:         xtile.insert %[[OUTPUT]] into %arg2[%c0] [1024] [1] : tensor<1024xf32> -> memref<1024xf32>
 )";
 
-  EXPECT_THAT(RunFileCheck(triton_mlir, kPattern),
+  EXPECT_THAT(RunFileCheck(triton_ir, kPattern),
               absl_testing::IsOkAndHolds(true))
-      << triton_mlir;
+      << triton_ir;
+}
+
+TEST_F(ExperimentalTilingTritonEmitterTest, Scan1DTiledEmitOk) {
+  constexpr absl::string_view kHloText = R"(
+scan_computation {
+  p_carry = f32[] parameter(0)
+  p_input = f32[] parameter(1)
+  add = f32[] add(p_carry, p_input)
+  ROOT tuple = (f32[], f32[]) tuple(add, add)
+}
+
+fusion_computation {
+  p0 = f32[1024]{0} parameter(0)
+  p1 = f32[] parameter(1)
+  scan = (f32[1024]{0}, f32[]) scan(p0, p1), dimensions={0}, num_carries=1, is_associative=true, to_apply=scan_computation, backend_config={"sizes": [256]}
+  ROOT gte = f32[1024]{0} get-tuple-element(scan), index=0
+}
+
+ENTRY main {
+  param0 = f32[1024]{0} parameter(0)
+  param1 = f32[] parameter(1)
+  ROOT triton_fusion = f32[1024]{0} fusion(param0, param1), kind=kCustom, calls=fusion_computation, backend_config={
+    "fusion_backend_config": {
+      "kind": "__triton_nested_gemm_fusion",
+      "block_level_fusion_config": {
+        "output_tiles": [{"sizes": []}],
+        "num_warps": 4,
+        "num_ctas": 1,
+        "num_stages": 1
+      }
+    }
+  }
+}
+)";
+
+  ASSERT_OK_AND_ASSIGN(std::string triton_ir,
+                       EmitAndCompileTritonFusion(kHloText));
+
+  // The carry of a tiled 1D scan is a rank-1 tensor, not a rank-zero tensor,
+  // which TritonGPU does not support. It is the last element of the scan
+  // output, selected with a masked reduction.
+  constexpr absl::string_view kPattern = R"(
+// CHECK-LABEL: @test_fn
+// CHECK:         %[[LOOP:.*]] = scf.for {{.*}} iter_args(%[[CARRY:.*]] = %{{.*}}) -> (tensor<1xf32>)
+// CHECK:           %[[BCAST_CARRY:.*]] = tt.broadcast %[[CARRY]] : tensor<1xf32> -> tensor<256xf32>
+// CHECK:           %[[OUTPUT:.*]] = arith.addf %[[BCAST_CARRY]], %{{.*}} : tensor<256xf32>
+// CHECK:           %[[MASKED:.*]] = arith.select %{{.*}}, %{{.*}}, %{{.*}} : tensor<256xi1>, tensor<256xi32>
+// CHECK:           %[[REDUCE:.*]] = "tt.reduce"(%[[MASKED]]) <{axis = 0 : i32}> ({
+// CHECK:             arith.ori
+// CHECK:           }) : (tensor<256xi32>) -> i32
+// CHECK:           %[[SPLAT:.*]] = tt.splat %[[REDUCE]] : i32 -> tensor<1xi32>
+// CHECK:           %[[NEXT_CARRY:.*]] = tt.bitcast %[[SPLAT]] : tensor<1xi32> -> tensor<1xf32>
+// CHECK-NOT:       tt.gather
+// CHECK:           scf.yield %[[NEXT_CARRY]] : tensor<1xf32>
+)";
+
+  EXPECT_THAT(RunFileCheck(triton_ir, kPattern),
+              absl_testing::IsOkAndHolds(true))
+      << triton_ir;
+}
+
+TEST_F(ExperimentalTilingTritonEmitterTest, Scan1DReverseTiledEmitOk) {
+  constexpr absl::string_view kHloText = R"(
+scan_computation {
+  p_carry = f32[] parameter(0)
+  p_input = f32[] parameter(1)
+  add = f32[] add(p_carry, p_input)
+  ROOT tuple = (f32[], f32[]) tuple(add, add)
+}
+
+fusion_computation {
+  p0 = f32[1024]{0} parameter(0)
+  p1 = f32[] parameter(1)
+  scan = (f32[1024]{0}, f32[]) scan(p0, p1), dimensions={0}, is_reverse=true, num_carries=1, is_associative=true, to_apply=scan_computation, backend_config={"sizes": [256]}
+  ROOT gte = f32[1024]{0} get-tuple-element(scan), index=0
+}
+
+ENTRY main {
+  param0 = f32[1024]{0} parameter(0)
+  param1 = f32[] parameter(1)
+  ROOT triton_fusion = f32[1024]{0} fusion(param0, param1), kind=kCustom, calls=fusion_computation, backend_config={
+    "fusion_backend_config": {
+      "kind": "__triton_nested_gemm_fusion",
+      "block_level_fusion_config": {
+        "output_tiles": [{"sizes": []}],
+        "num_warps": 4,
+        "num_ctas": 1,
+        "num_stages": 1
+      }
+    }
+  }
+}
+)";
+
+  ASSERT_OK_AND_ASSIGN(std::string triton_ir,
+                       EmitAndCompileTritonFusion(kHloText));
+
+  // In a reverse tiled scan, the carry is element 0 of the scan output.
+  constexpr absl::string_view kPattern = R"(
+// CHECK-LABEL: @test_fn
+// CHECK:         %[[LOOP:.*]] = scf.for {{.*}} iter_args(%[[CARRY:.*]] = %{{.*}}) -> (tensor<1xf32>)
+// CHECK:           %[[BCAST_CARRY:.*]] = tt.broadcast %[[CARRY]] : tensor<1xf32> -> tensor<256xf32>
+// CHECK:           %[[OUTPUT:.*]] = arith.addf %[[BCAST_CARRY]], %{{.*}} : tensor<256xf32>
+// CHECK:           %[[MASKED:.*]] = arith.select %{{.*}}, %{{.*}}, %{{.*}} : tensor<256xi1>, tensor<256xi32>
+// CHECK:           %[[REDUCE:.*]] = "tt.reduce"(%[[MASKED]]) <{axis = 0 : i32}> ({
+// CHECK:             arith.ori
+// CHECK:           }) : (tensor<256xi32>) -> i32
+// CHECK:           %[[SPLAT:.*]] = tt.splat %[[REDUCE]] : i32 -> tensor<1xi32>
+// CHECK:           %[[NEXT_CARRY:.*]] = tt.bitcast %[[SPLAT]] : tensor<1xi32> -> tensor<1xf32>
+// CHECK-NOT:       tt.gather
+// CHECK:           scf.yield %[[NEXT_CARRY]] : tensor<1xf32>
+)";
+
+  EXPECT_THAT(RunFileCheck(triton_ir, kPattern),
+              absl_testing::IsOkAndHolds(true))
+      << triton_ir;
+}
+
+// The carry is extracted by bitcasting the scan output to an integer of the
+// same width, which is i16 for bf16. Compiling the fusion covers lowering the
+// `or` reduction over i16, which cannot use the 32-bit `redux.sync` fast path.
+TEST_F(ExperimentalTilingTritonEmitterTest, ScanBf16TiledEmitOk) {
+  constexpr absl::string_view kHloText = R"(
+scan_computation {
+  p_carry = bf16[] parameter(0)
+  p_input = bf16[] parameter(1)
+  add = bf16[] add(p_carry, p_input)
+  ROOT tuple = (bf16[], bf16[]) tuple(add, add)
+}
+
+fusion_computation {
+  p0 = bf16[1024]{0} parameter(0)
+  p1 = bf16[] parameter(1)
+  scan = (bf16[1024]{0}, bf16[]) scan(p0, p1), dimensions={0}, num_carries=1, is_associative=true, to_apply=scan_computation, backend_config={"sizes": [256]}
+  ROOT gte = bf16[1024]{0} get-tuple-element(scan), index=0
+}
+
+ENTRY main {
+  param0 = bf16[1024]{0} parameter(0)
+  param1 = bf16[] parameter(1)
+  ROOT triton_fusion = bf16[1024]{0} fusion(param0, param1), kind=kCustom, calls=fusion_computation, backend_config={
+    "fusion_backend_config": {
+      "kind": "__triton_nested_gemm_fusion",
+      "block_level_fusion_config": {
+        "output_tiles": [{"sizes": []}],
+        "num_warps": 4,
+        "num_ctas": 1,
+        "num_stages": 1
+      }
+    }
+  }
+}
+)";
+
+  ASSERT_OK_AND_ASSIGN(std::string triton_ir,
+                       EmitAndCompileTritonFusion(kHloText));
+
+  constexpr absl::string_view kPattern = R"(
+// CHECK-LABEL: @test_fn
+// CHECK:         %[[LOOP:.*]] = scf.for {{.*}} iter_args(%[[CARRY:.*]] = %{{.*}}) -> (tensor<1xbf16>)
+// CHECK:           %[[MASKED:.*]] = arith.select %{{.*}}, %{{.*}}, %{{.*}} : tensor<256xi1>, tensor<256xi16>
+// CHECK:           %[[REDUCE:.*]] = "tt.reduce"(%[[MASKED]]) <{axis = 0 : i32}> ({
+// CHECK:             arith.ori
+// CHECK:           }) : (tensor<256xi16>) -> i16
+// CHECK:           %[[SPLAT:.*]] = tt.splat %[[REDUCE]] : i16 -> tensor<1xi16>
+// CHECK:           %[[NEXT_CARRY:.*]] = tt.bitcast %[[SPLAT]] : tensor<1xi16> -> tensor<1xbf16>
+// CHECK-NOT:       tt.gather
+// CHECK:           scf.yield %[[NEXT_CARRY]] : tensor<1xbf16>
+)";
+
+  EXPECT_THAT(RunFileCheck(triton_ir, kPattern),
+              absl_testing::IsOkAndHolds(true))
+      << triton_ir;
+}
+
+TEST_F(ExperimentalTilingTritonEmitterTest, Scan2DTiledEmitOk) {
+  constexpr absl::string_view kHloText = R"(
+scan_computation {
+  p_carry = f32[16]{0} parameter(0)
+  p_input = f32[16]{0} parameter(1)
+  add = f32[16]{0} add(p_carry, p_input)
+  ROOT tuple = (f32[16]{0}, f32[16]{0}) tuple(add, add)
+}
+
+fusion_computation {
+  p0 = f32[16,1024]{1,0} parameter(0)
+  p1 = f32[16]{0} parameter(1)
+  scan = (f32[16,1024]{1,0}, f32[16]{0}) scan(p0, p1), dimensions={1}, num_carries=1, is_associative=true, to_apply=scan_computation, backend_config={"sizes": [256]}
+  ROOT gte = f32[16,1024]{1,0} get-tuple-element(scan), index=0
+}
+
+ENTRY main {
+  param0 = f32[16,1024]{1,0} parameter(0)
+  param1 = f32[16]{0} parameter(1)
+  ROOT triton_fusion = f32[16,1024]{1,0} fusion(param0, param1), kind=kCustom, calls=fusion_computation, backend_config={
+    "fusion_backend_config": {
+      "kind": "__triton_nested_gemm_fusion",
+      "block_level_fusion_config": {
+        "output_tiles": [{"sizes": [16]}],
+        "num_warps": 4,
+        "num_ctas": 1,
+        "num_stages": 1
+      }
+    }
+  }
+}
+)";
+
+  ASSERT_OK_AND_ASSIGN(std::string triton_ir,
+                       EmitAndCompileTritonFusion(kHloText));
+
+  constexpr absl::string_view kPattern = R"(
+// CHECK-LABEL: @test_fn
+// CHECK:         %[[LOOP:.*]] = scf.for {{.*}} iter_args(%[[CARRY:.*]] = %{{.*}}) -> (tensor<16x1xf32>)
+// CHECK:           %[[BCAST_CARRY:.*]] = tt.broadcast %[[CARRY]] : tensor<16x1xf32> -> tensor<16x256xf32>
+// CHECK:           %[[OUTPUT:.*]] = arith.addf %[[BCAST_CARRY]], %{{.*}} : tensor<16x256xf32>
+// CHECK:           %[[MASKED:.*]] = arith.select %{{.*}}, %{{.*}}, %{{.*}} : tensor<16x256xi1>, tensor<16x256xi32>
+// CHECK:           %[[REDUCE:.*]] = "tt.reduce"(%[[MASKED]]) <{axis = 1 : i32}> ({
+// CHECK:             arith.ori
+// CHECK:           }) : (tensor<16x256xi32>) -> tensor<16xi32>
+// CHECK:           %[[EXPAND:.*]] = tt.expand_dims %[[REDUCE]] {axis = 1 : i32} : tensor<16xi32> -> tensor<16x1xi32>
+// CHECK:           %[[NEXT_CARRY:.*]] = tt.bitcast %[[EXPAND]] : tensor<16x1xi32> -> tensor<16x1xf32>
+// CHECK-NOT:       tt.gather
+// CHECK:           scf.yield %[[NEXT_CARRY]] : tensor<16x1xf32>
+)";
+
+  EXPECT_THAT(RunFileCheck(triton_ir, kPattern),
+              absl_testing::IsOkAndHolds(true))
+      << triton_ir;
 }
 
 class UnsignedIntegerOpsTest
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_shared_dialect_test.cc b/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_shared_dialect_test.cc
index 5972e60d07a6cf..ce1a4f839e0350 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_shared_dialect_test.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_shared_dialect_test.cc
@@ -290,9 +290,10 @@ ENTRY e {
 // CHECK-LABEL: @xtile_dialect_fn
 // CHECK:         %[[EXTRACT0:.*]] = xtile.extract %arg0[%{{.*}}] [1024] [1] : memref<1024xf32> -> tensor<1024xf32>
 // CHECK:         %[[EXTRACT1:.*]] = xtile.extract %arg1[] [] [] : memref<f32> -> tensor<f32>
-// CHECK:         %[[OUTPUT:.*]], %{{.*}} = xtile.scan(%[[EXTRACT0]]) inits(%[[EXTRACT1]])
+// CHECK:         %[[INIT:.*]] = stablehlo.reshape %[[EXTRACT1]] : (tensor<f32>) -> tensor<1xf32>
+// CHECK:         %[[OUTPUT:.*]], %{{.*}} = xtile.scan(%[[EXTRACT0]]) inits(%[[INIT]])
 // CHECK-SAME:        dimension = 0 {scan_dim_size = 1024 : i64}
-// CHECK-SAME:        : (tensor<1024xf32>), (tensor<f32>) -> (tensor<1024xf32>), (tensor<f32>) {
+// CHECK-SAME:        : (tensor<1024xf32>), (tensor<1xf32>) -> (tensor<1024xf32>), (tensor<1xf32>) {
 // CHECK:         ^bb0(%[[INPUT:.*]]: tensor<f32>, %[[CARRY:.*]]: tensor<f32>):
 // CHECK:           %[[ADD:.*]] = stablehlo.add %[[INPUT]], %[[CARRY]] : tensor<f32>
 // CHECK:           stablehlo.return %[[ADD]], %[[ADD]] : tensor<f32>, tensor<f32>
@@ -635,6 +636,54 @@ ENTRY e {
   )"));
 }
 
+TEST_F(XTileDialectTest, HloReverseScanWithLoop) {
+  constexpr absl::string_view kHloText = R"(
+HloModule t
+
+add_computation {
+  p0 = f32[] parameter(0)
+  p1 = f32[] parameter(1)
+  add = f32[] add(p0, p1)
+  ROOT tuple = (f32[], f32[]) tuple(add, add)
+}
+
+scan_fusion {
+  p0 = f32[1024] parameter(0)
+  p1 = f32[] parameter(1)
+  scan = (f32[1024], f32[]) scan(p0, p1),
+    dimensions={0}, is_reverse=true, num_carries=1, is_associative=true,
+    to_apply=add_computation, backend_config={sizes:[128]}
+  ROOT get-tuple-element = f32[1024] get-tuple-element(scan), index=0
+}
+
+ENTRY e {
+  p0 = f32[1024] parameter(0)
+  p1 = f32[] parameter(1)
+  ROOT custom-call = f32[1024] fusion(p0, p1), kind=kCustom,
+    calls=scan_fusion,
+    backend_config={"fusion_backend_config": {kind: "__triton"}}
+})";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnUnverifiedModule(kHloText));
+  module->mutable_config()
+      .mutable_debug_options()
+      .set_xla_gpu_experimental_enable_tiling_propagation(true);
+
+  BlockLevelParameters block_level_parameters;
+  block_level_parameters.output_tile_sizes = {{}};
+
+  EXPECT_OK(CreateXTileIrAndFileCheck(
+      *module->GetComputationWithName("scan_fusion"), block_level_parameters,
+      R"(
+    CHECK: scf.for %[[IV:.*]] = %{{.*}} to %{{.*}} step %{{.*}} iter_args(%[[CARRY:.*]] = %{{.*}})
+    CHECK: %[[REVERSE_IDX:.*]] = arith.subi %{{.*}}, %[[IV]] : index
+    CHECK: %[[INPUT_TILE:.*]] = xtile.extract %arg0
+    CHECK: %[[SCAN_OUT:.*]], %[[NEW_CARRY:.*]] = xtile.scan(%[[INPUT_TILE]]) inits(%[[CARRY]]) dimension = 0 {is_reverse = true
+    CHECK: xtile.insert %[[SCAN_OUT]] into %arg2
+    CHECK: scf.yield %[[NEW_CARRY]]
+  )"));
+}
+
 // When the scan dimension is not divisible by the tile size, the input tile of
 // the final loop iteration extends past the global dimension bound. These
 // out-of-bounds elements must be dynamically masked (using the loop induction
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/passes.td b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/passes.td
index 3c6e755c1c18b7..bad02e24865796 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/passes.td
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/passes.td
@@ -243,6 +243,7 @@ def XTileLowerToTritonPass
     : Pass<"xtile-lower-to-triton", "mlir::ModuleOp"> {
   let summary = "Lowers XTile operations to their Triton equivalent.";
   let dependentDialects = [
+    "::mlir::arith::ArithDialect",
     "::mlir::stablehlo::StablehloDialect",
     "::mlir::triton::TritonDialect",
     "::xla::xtile::XTileDialect",
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_squeeze_dims.mlir b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_squeeze_dims.mlir
index 31b28dbdab386d..63d0cbc0faa245 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_squeeze_dims.mlir
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_squeeze_dims.mlir
@@ -135,6 +135,17 @@ func.func @skip_reshape_with_attr(%arg0: tensor<4x1xf32>) -> tensor<4xf32> {
 
 // -----
 
+// A rank zero result has no unit dimensions to squeeze.
+// CHECK-LABEL: func @skip_reshape_to_rank_zero
+func.func @skip_reshape_to_rank_zero(%arg0: tensor<1xf32>) -> tensor<f32> {
+  // CHECK-NOT: triton_xla.squeeze_dims
+  // CHECK: tt.reshape {{.*}} : tensor<1xf32> -> tensor<f32>
+  %0 = tt.reshape %arg0 : tensor<1xf32> -> tensor<f32>
+  return %0 : tensor<f32>
+}
+
+// -----
+
 #arg_enc = #ttg.blocked<{sizePerThread = [1, 1], threadsPerWarp = [1, 32], warpsPerCTA = [1, 1], order = [1, 0]}>
 #res_enc = #ttg.blocked<{sizePerThread = [1], threadsPerWarp = [32], warpsPerCTA = [1], order = [0]}>
 module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32} {
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/xtile_to_triton_lowering.mlir b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/xtile_to_triton_lowering.mlir
index 5c180c00f7cd8a..1657db56a13535 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/xtile_to_triton_lowering.mlir
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/xtile_to_triton_lowering.mlir
@@ -80,19 +80,19 @@ func.func @lower_dot_scaled_in_loop_non_canonical(
 // -----
 
 // CHECK-LABEL: func.func @scan_lowering(
-// CHECK-SAME: %[[INPUT0:.*]]: tensor<16x16x16xf32>, %[[INIT0:.*]]: tensor<16x16xf32>
-func.func @scan_lowering(%input0: tensor<16x16x16xf32>, %init0: tensor<16x16xf32>) -> (tensor<16x16x16xf32>) {
+// CHECK-SAME: %[[INPUT0:.*]]: tensor<16x16x16xf32>, %[[INIT0:.*]]: tensor<16x16x1xf32>
+func.func @scan_lowering(%input0: tensor<16x16x16xf32>, %init0: tensor<16x16x1xf32>) -> (tensor<16x16x16xf32>) {
   // CHECK: %[[SCAN:.*]] = "tt.scan"(%[[INPUT0]]) <{axis = 2 : i32, reverse = false}> ({
   // CHECK:   ^bb0(%[[ARG0:.*]]: f32, %[[ARG1:.*]]: f32):
   // CHECK:     %[[ADD:.*]] = arith.addf %[[ARG0]], %[[ARG1]] : f32
   // CHECK:     tt.scan.return %[[ADD]] : f32
   // CHECK: }
 
-  // CHECK-DAG: %[[BCAST_INIT0:.*]] = stablehlo.broadcast_in_dim %[[INIT0]], dims = [0, 1] : (tensor<16x16xf32>) -> tensor<16x16x16xf32>
+  // CHECK-DAG: %[[BCAST_INIT0:.*]] = tt.broadcast %[[INIT0]] : tensor<16x16x1xf32> -> tensor<16x16x16xf32>
   // CHECK-DAG: %[[RES1:.*]] = arith.addf %[[BCAST_INIT0]], %[[SCAN]] : tensor<16x16x16xf32>
   // CHECK: return %[[RES1]] : tensor<16x16x16xf32>
 
-  %0, %1 = xtile.scan(%input0) inits(%init0) dimension = 2 {scan_dim_size = 16 : i64} : (tensor<16x16x16xf32>), (tensor<16x16xf32>) -> (tensor<16x16x16xf32>), (tensor<16x16xf32>) {
+  %0, %1 = xtile.scan(%input0) inits(%init0) dimension = 2 {scan_dim_size = 16 : i64} : (tensor<16x16x16xf32>), (tensor<16x16x1xf32>) -> (tensor<16x16x16xf32>), (tensor<16x16x1xf32>) {
   ^bb0(%arg0: f32, %arg1: f32):
     %add = arith.addf %arg0, %arg1 : f32
     xtile.yield %add : f32
@@ -100,6 +100,114 @@ func.func @scan_lowering(%input0: tensor<16x16x16xf32>, %init0: tensor<16x16xf32
   return %0 : tensor<16x16x16xf32>
 }
 
+// -----
+
+// CHECK-LABEL: func.func @scan_lowering_carry(
+// CHECK-SAME: %[[INPUT0:.*]]: tensor<16x16x16xf32>, %[[INIT0:.*]]: tensor<16x16x1xf32>
+func.func @scan_lowering_carry(%input0: tensor<16x16x16xf32>, %init0: tensor<16x16x1xf32>) -> (tensor<16x16x16xf32>, tensor<16x16x1xf32>) {
+  // CHECK-DAG: %[[ZERO:.*]] = arith.constant dense<0> : tensor<16x16x16xi32>
+  // CHECK-DAG: %[[LAST:.*]] = arith.constant dense<15> : tensor<16x16x16xi32>
+  // CHECK: %[[SCAN:.*]] = "tt.scan"(%[[INPUT0]]) <{axis = 2 : i32, reverse = false}> ({
+  // CHECK: %[[BCAST_INIT0:.*]] = tt.broadcast %[[INIT0]] : tensor<16x16x1xf32> -> tensor<16x16x16xf32>
+  // CHECK: %[[RES0:.*]] = arith.addf %[[BCAST_INIT0]], %[[SCAN]] : tensor<16x16x16xf32>
+
+  // The carry is the last element along the scan dimension, which is selected
+  // with a masked reduction over the bitcast integer values.
+  // CHECK: %[[RANGE:.*]] = tt.make_range {end = 16 : i32, start = 0 : i32} : tensor<16xi32>
+  // CHECK: %[[EXPAND0:.*]] = tt.expand_dims %[[RANGE]] {axis = 0 : i32} : tensor<16xi32> -> tensor<1x16xi32>
+  // CHECK: %[[EXPAND1:.*]] = tt.expand_dims %[[EXPAND0]] {axis = 1 : i32} : tensor<1x16xi32> -> tensor<1x1x16xi32>
+  // CHECK: %[[IOTA:.*]] = tt.broadcast %[[EXPAND1]] : tensor<1x1x16xi32> -> tensor<16x16x16xi32>
+  // CHECK: %[[MASK:.*]] = arith.cmpi eq, %[[IOTA]], %[[LAST]] : tensor<16x16x16xi32>
+  // CHECK: %[[BITCAST:.*]] = tt.bitcast %[[RES0]] : tensor<16x16x16xf32> -> tensor<16x16x16xi32>
+  // CHECK: %[[MASKED:.*]] = arith.select %[[MASK]], %[[BITCAST]], %[[ZERO]]
+  // CHECK: %[[REDUCE:.*]] = "tt.reduce"(%[[MASKED]]) <{axis = 2 : i32}> ({
+  // CHECK:   ^bb0(%[[LHS:.*]]: i32, %[[RHS:.*]]: i32):
+  // CHECK:     %[[OR:.*]] = arith.ori %[[LHS]], %[[RHS]] : i32
+  // CHECK:     tt.reduce.return %[[OR]] : i32
+  // CHECK: }) : (tensor<16x16x16xi32>) -> tensor<16x16xi32>
+  // CHECK: %[[EXPAND:.*]] = tt.expand_dims %[[REDUCE]] {axis = 2 : i32} : tensor<16x16xi32> -> tensor<16x16x1xi32>
+  // CHECK: %[[RES1:.*]] = tt.bitcast %[[EXPAND]] : tensor<16x16x1xi32> -> tensor<16x16x1xf32>
+  // CHECK-NOT: tt.gather
+  // CHECK-NOT: tt.reshape
+  // CHECK: return %[[RES0]], %[[RES1]]
+
+  %0, %1 = xtile.scan(%input0) inits(%init0) dimension = 2 {scan_dim_size = 16 : i64} : (tensor<16x16x16xf32>), (tensor<16x16x1xf32>) -> (tensor<16x16x16xf32>), (tensor<16x16x1xf32>) {
+  ^bb0(%arg0: f32, %arg1: f32):
+    %add = arith.addf %arg0, %arg1 : f32
+    xtile.yield %add : f32
+  }
+  return %0, %1 : tensor<16x16x16xf32>, tensor<16x16x1xf32>
+}
+
+// -----
+
+// The carry of a 1D scan is a rank-1 tensor, not a rank-zero tensor, which
+// TritonGPU does not support.
+// CHECK-LABEL: func.func @scan_lowering_1d_carry(
+// CHECK-SAME: %[[INPUT0:.*]]: tensor<16xf32>, %[[INIT0:.*]]: tensor<1xf32>
+func.func @scan_lowering_1d_carry(%input0: tensor<16xf32>, %init0: tensor<1xf32>) -> (tensor<16xf32>, tensor<1xf32>) {
+  // CHECK-DAG: %[[ZERO:.*]] = arith.constant dense<0> : tensor<16xi32>
+  // CHECK-DAG: %[[LAST:.*]] = arith.constant dense<15> : tensor<16xi32>
+  // CHECK: %[[SCAN:.*]] = "tt.scan"(%[[INPUT0]]) <{axis = 0 : i32, reverse = false}> ({
+  // CHECK: %[[BCAST_INIT0:.*]] = tt.broadcast %[[INIT0]] : tensor<1xf32> -> tensor<16xf32>
+  // CHECK: %[[RES0:.*]] = arith.addf %[[BCAST_INIT0]], %[[SCAN]] : tensor<16xf32>
+  // CHECK: %[[RANGE:.*]] = tt.make_range {end = 16 : i32, start = 0 : i32} : tensor<16xi32>
+  // CHECK: %[[MASK:.*]] = arith.cmpi eq, %[[RANGE]], %[[LAST]] : tensor<16xi32>
+  // CHECK: %[[BITCAST:.*]] = tt.bitcast %[[RES0]] : tensor<16xf32> -> tensor<16xi32>
+  // CHECK: %[[MASKED:.*]] = arith.select %[[MASK]], %[[BITCAST]], %[[ZERO]]
+  // CHECK: %[[REDUCE:.*]] = "tt.reduce"(%[[MASKED]]) <{axis = 0 : i32}> ({
+  // CHECK:   ^bb0(%[[LHS:.*]]: i32, %[[RHS:.*]]: i32):
+  // CHECK:     %[[OR:.*]] = arith.ori %[[LHS]], %[[RHS]] : i32
+  // CHECK:     tt.reduce.return %[[OR]] : i32
+  // CHECK: }) : (tensor<16xi32>) -> i32
+  // Reducing a rank-1 tensor returns a scalar, which is splatted back to the
+  // unit dimension of the carry.
+  // CHECK: %[[SPLAT:.*]] = tt.splat %[[REDUCE]] : i32 -> tensor<1xi32>
+  // CHECK: %[[RES1:.*]] = tt.bitcast %[[SPLAT]] : tensor<1xi32> -> tensor<1xf32>
+  // CHECK-NOT: tt.gather
+  // CHECK-NOT: tt.reshape
+  // CHECK: return %[[RES0]], %[[RES1]]
+  %0, %1 = xtile.scan(%input0) inits(%init0) dimension = 0 {scan_dim_size = 16 : i64} : (tensor<16xf32>), (tensor<1xf32>) -> (tensor<16xf32>), (tensor<1xf32>) {
+  ^bb0(%arg0: f32, %arg1: f32):
+    %add = arith.addf %arg0, %arg1 : f32
+    xtile.yield %add : f32
+  }
+  return %0, %1 : tensor<16xf32>, tensor<1xf32>
+}
+
+// -----
+
+// The carry of a reverse scan is the first element along the scan dimension.
+// CHECK-LABEL: func.func @reverse_scan_lowering_carry(
+func.func @reverse_scan_lowering_carry(%input0: tensor<16xf32>, %init0: tensor<1xf32>) -> (tensor<16xf32>, tensor<1xf32>) {
+  // CHECK-DAG: %[[ZERO:.*]] = arith.constant dense<0> : tensor<16xi32>
+  // CHECK: %[[SCAN:.*]] = "tt.scan"(%{{.*}}) <{axis = 0 : i32, reverse = true}> ({
+  // CHECK: %[[RES0:.*]] = arith.addf %{{.*}}, %[[SCAN]] : tensor<16xf32>
+  // The mask therefore compares against index 0 instead of the last index.
+  // CHECK: %[[RANGE:.*]] = tt.make_range {end = 16 : i32, start = 0 : i32} : tensor<16xi32>
+  // CHECK: %[[MASK:.*]] = arith.cmpi eq, %[[RANGE]], %[[ZERO]] : tensor<16xi32>
+  // CHECK: %[[BITCAST:.*]] = tt.bitcast %[[RES0]] : tensor<16xf32> -> tensor<16xi32>
+  // CHECK: %[[MASKED:.*]] = arith.select %[[MASK]], %[[BITCAST]], %[[ZERO]]
+  // CHECK: %[[REDUCE:.*]] = "tt.reduce"(%[[MASKED]]) <{axis = 0 : i32}> ({
+  // CHECK:   ^bb0(%[[LHS:.*]]: i32, %[[RHS:.*]]: i32):
+  // CHECK:     %[[OR:.*]] = arith.ori %[[LHS]], %[[RHS]] : i32
+  // CHECK:     tt.reduce.return %[[OR]] : i32
+  // CHECK: }) : (tensor<16xi32>) -> i32
+  // CHECK: %[[SPLAT:.*]] = tt.splat %[[REDUCE]] : i32 -> tensor<1xi32>
+  // CHECK: %[[RES1:.*]] = tt.bitcast %[[SPLAT]] : tensor<1xi32> -> tensor<1xf32>
+  // CHECK-NOT: tt.gather
+  // CHECK-NOT: tt.reshape
+  // CHECK: return %[[RES0]], %[[RES1]]
+  %0, %1 = xtile.scan(%input0) inits(%init0) dimension = 0 {scan_dim_size = 16 : i64, is_reverse = true} : (tensor<16xf32>), (tensor<1xf32>) -> (tensor<16xf32>), (tensor<1xf32>) {
+  ^bb0(%arg0: f32, %arg1: f32):
+    %add = arith.addf %arg0, %arg1 : f32
+    xtile.yield %add : f32
+  }
+  return %0, %1 : tensor<16xf32>, tensor<1xf32>
+}
+
+// -----
+
 // CHECK: func @lower_dot_scaled_without_add_falls_back_to_xtile(%[[LHS:.*]]: tensor<128x128xf8E5M2>, %[[LHS_SCALE:.*]]: tensor<128x4xi8>, %[[RHS:.*]]: tensor<128x256xf8E5M2>, %[[RHS_SCALE:.*]]: tensor<256x4xi8>) -> tensor<128x256xf32> {
 func.func @lower_dot_scaled_without_add_falls_back_to_xtile(
   %lhs: tensor<128x128xf8E5M2>, %lhs_scale: tensor<128x4xi8>,
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/triton_xla_squeeze_dims_pass.cc b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/triton_xla_squeeze_dims_pass.cc
index 355c1c9f401e51..0aae6164de3965 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/triton_xla_squeeze_dims_pass.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/triton_xla_squeeze_dims_pass.cc
@@ -62,7 +62,7 @@ SmallVector<uint32_t> GetDimsToSqueeze(RankedTensorType type) {
       result.push_back(dim);
     }
   }
-  if (result.size() == type.getRank()) {
+  if (!result.empty() && result.size() == type.getRank()) {
     result.pop_back();  // Keep one unit dimension.
   }
   return result;
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/xtile_lower_to_triton.cc b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/xtile_lower_to_triton.cc
index ade8d93e5be0c1..7c6b01150f1370 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/xtile_lower_to_triton.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/xtile_lower_to_triton.cc
@@ -22,6 +22,8 @@ limitations under the License.
 #include "absl/strings/str_cat.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVectorExtras.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Arith/IR/ArithAttributes.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/Builders.h"
@@ -329,25 +331,15 @@ class LowerScan : public mlir::OpRewritePattern<::xla::xtile::ScanOp> {
   static SmallVector<Value> FoldInitValues(::xla::xtile::ScanOp op,
                                            ttir::ScanOp triton_scan_op,
                                            mlir::PatternRewriter& rewriter) {
-    int32_t axis = op.getDimension();
     SmallVector<Value> init_and_results;
     init_and_results.reserve(triton_scan_op.getNumResults() * 2);
 
+    // The init only has a unit dimension along the scan dimension and is
+    // combined with every element of the scan result.
     for (auto [result, init_val] :
          llvm::zip_equal(triton_scan_op.getResults(), op.getInits())) {
-      auto result_type = mlir::cast<mlir::RankedTensorType>(result.getType());
-
-      SmallVector<int64_t> bcast_dims;
-      bcast_dims.reserve(result_type.getRank() - 1);
-      for (int64_t d = 0; d < result_type.getRank(); ++d) {
-        if (d != axis) {
-          bcast_dims.push_back(d);
-        }
-      }
-
-      init_and_results.push_back(mlir::stablehlo::BroadcastInDimOp::create(
-          rewriter, op.getLoc(), result_type, init_val,
-          rewriter.getDenseI64ArrayAttr(bcast_dims)));
+      init_and_results.push_back(ttir::BroadcastOp::create(
+          rewriter, op.getLoc(), result.getType(), init_val));
     }
 
     llvm::append_range(init_and_results, triton_scan_op.getResults());
@@ -363,42 +355,115 @@ class LowerScan : public mlir::OpRewritePattern<::xla::xtile::ScanOp> {
     return tensor_outputs;
   }
 
-  static SmallVector<Value> ExtractCarries(::xla::xtile::ScanOp op,
-                                           ArrayRef<Value> outputs,
-                                           mlir::PatternRewriter& rewriter) {
-    int32_t axis = op.getDimension();
-    bool reverse = op.getIsReverse();
-
-    SmallVector<Value> carries;
-    carries.reserve(op.getCarries().size());
-    for (auto [output, carry] : llvm::zip(outputs, op.getCarries())) {
-      auto result_type = mlir::cast<mlir::RankedTensorType>(output.getType());
-      SmallVector<OpFoldResult> offsets, sizes, strides;
-      offsets.reserve(result_type.getRank());
-      sizes.reserve(result_type.getRank());
-      strides.reserve(result_type.getRank());
-
-      for (int64_t d = 0; d < result_type.getRank(); ++d) {
-        if (d == axis) {
-          int64_t start_idx = reverse ? 0 : (result_type.getDimSize(d) - 1);
-          offsets.push_back(rewriter.getIndexAttr(start_idx));
-          sizes.push_back(rewriter.getIndexAttr(1));
-        } else {
-          offsets.push_back(rewriter.getIndexAttr(0));
-          sizes.push_back(rewriter.getIndexAttr(result_type.getDimSize(d)));
-        }
-        strides.push_back(rewriter.getIndexAttr(1));
+  // Returns a tensor of the same shape as `type` with the index along `axis`.
+  static Value CreateIota(mlir::RankedTensorType type, int32_t axis,
+                          mlir::Location loc, mlir::PatternRewriter& rewriter) {
+    mlir::Type i32_type = rewriter.getI32Type();
+    int64_t axis_size = type.getDimSize(axis);
+    Value iota = ttir::MakeRangeOp::create(
+        rewriter, loc, mlir::RankedTensorType::get({axis_size}, i32_type),
+        /*start=*/0, /*end=*/axis_size);
+    // Inserting the unit dimensions in increasing order leaves the range in
+    // dimension `axis`.
+    for (int32_t dim = 0; dim < type.getRank(); ++dim) {
+      if (dim != axis) {
+        iota = ttir::ExpandDimsOp::create(rewriter, loc, iota, dim);
       }
-      auto carry_type = mlir::cast<mlir::RankedTensorType>(carry.getType());
+    }
+    return ttir::BroadcastOp::create(rewriter, loc, type.clone(i32_type), iota);
+  }
 
-      carries.push_back(mlir::tensor::ExtractSliceOp::create(
-          rewriter, op.getLoc(), carry_type, output, offsets, sizes, strides));
+  // Returns the last element that the scan produces along `axis`, which is the
+  // element at index 0 for reverse scans and at the last index otherwise,
+  // keeping `axis` as a unit dimension.
+  //
+  // `tt.scan` does not return the final carry, and extracting it with a
+  // `tt.gather` would write the entire tile to shared memory (see
+  // GatherLoweringHelper::getScratchSizeInBytes()). Instead, mask all other
+  // elements to zero and combine them with an `or` reduction over the bitcast
+  // integer values. Zero is the neutral element of `or`, so the reduction
+  // returns the single unmasked element, and because `or` is idempotent it does
+  // so no matter in which order and how often a layout combines the elements.
+  // PyTorch inductor uses the same approach (see triton_helpers.select_one).
+  static Value ExtractCarry(Value value, int32_t axis, bool is_reverse,
+                            mlir::Location loc,
+                            mlir::PatternRewriter& rewriter) {
+    auto type = mlir::cast<mlir::RankedTensorType>(value.getType());
+    mlir::Type element_type = type.getElementType();
+    mlir::Type int_type =
+        rewriter.getIntegerType(element_type.getIntOrFloatBitWidth());
+
+    int32_t index =
+        is_reverse ? 0 : static_cast<int32_t>(type.getDimSize(axis) - 1);
+    Value iota = CreateIota(type, axis, loc, rewriter);
+    Value indices = mlir::arith::ConstantOp::create(
+        rewriter, loc,
+        mlir::DenseElementsAttr::get(
+            mlir::cast<mlir::ShapedType>(iota.getType()), index));
+    Value mask = mlir::arith::CmpIOp::create(
+        rewriter, loc, mlir::arith::CmpIPredicate::eq, iota, indices);
+
+    mlir::RankedTensorType int_tensor_type = type.clone(int_type);
+    if (int_type != element_type) {
+      value = ttir::BitcastOp::create(rewriter, loc, int_tensor_type, value);
     }
-    return carries;
+    Value zero = mlir::arith::ConstantOp::create(
+        rewriter, loc, rewriter.getZeroAttr(int_tensor_type));
+    Value masked =
+        mlir::arith::SelectOp::create(rewriter, loc, mask, value, zero);
+
+    // `tt.reduce` drops the reduced dimension and returns a scalar for rank-1
+    // inputs.
+    SmallVector<int64_t> reduced_shape(type.getShape());
+    reduced_shape.erase(reduced_shape.begin() + axis);
+    mlir::Type reduced_type =
+        reduced_shape.empty()
+            ? int_type
+            : mlir::RankedTensorType::get(reduced_shape, int_type);
+    auto reduce_op =
+        ttir::ReduceOp::create(rewriter, loc, reduced_type, masked, axis);
+    {
+      mlir::OpBuilder::InsertionGuard guard(rewriter);
+      mlir::Block* block = rewriter.createBlock(
+          &reduce_op.getCombineOp(), reduce_op.getCombineOp().begin(),
+          {int_type, int_type}, {loc, loc});
+      Value combined = mlir::arith::OrIOp::create(
+          rewriter, loc, block->getArgument(0), block->getArgument(1));
+      ttir::ReduceReturnOp::create(rewriter, loc, combined);
+    }
+
+    // Restore the scan dimension as a unit dimension.
+    SmallVector<int64_t> carry_shape(type.getShape());
+    carry_shape[axis] = 1;
+    Value carry = reduce_op.getResult().front();
+    if (reduced_shape.empty()) {
+      carry = ttir::SplatOp::create(
+          rewriter, loc, mlir::RankedTensorType::get(carry_shape, int_type),
+          carry);
+    } else {
+      carry = ttir::ExpandDimsOp::create(rewriter, loc, carry, axis);
+    }
+    if (int_type != element_type) {
+      carry = ttir::BitcastOp::create(
+          rewriter, loc, mlir::RankedTensorType::get(carry_shape, element_type),
+          carry);
+    }
+    return carry;
+  }
+
+  static SmallVector<Value> ExtractCarries(::xla::xtile::ScanOp op,
+                                           ArrayRef<Value> outputs,
+                                           mlir::PatternRewriter& rewriter) {
+    return llvm::map_to_vector(outputs, [&](Value output) {
+      return ExtractCarry(output, op.getDimension(), op.getIsReverse(),
+                          op.getLoc(), rewriter);
+    });
   }
 
   mlir::LogicalResult matchAndRewrite(
       ::xla::xtile::ScanOp op, mlir::PatternRewriter& rewriter) const override {
+    // `xtile.scan` verifies that the inits and carries keep the scan dimension
+    // as a unit dimension, which the lowering below relies on.
     ttir::ScanOp triton_scan_op = CreateTritonScan(op, rewriter);
 
     rewriter.setInsertionPointAfter(triton_scan_op);
diff --git a/third_party/xla/xla/backends/profiler/gpu/BUILD b/third_party/xla/xla/backends/profiler/gpu/BUILD
index 322540722c7580..8208e5ad72e141 100644
--- a/third_party/xla/xla/backends/profiler/gpu/BUILD
+++ b/third_party/xla/xla/backends/profiler/gpu/BUILD
@@ -606,17 +606,12 @@ cc_library(
     ],
 )
 
-xla_cc_test(
+xla_test(
     name = "rocm_tracer_test",
     size = "small",
     srcs = ["rocm_tracer_test.cc"],
-    tags = [
-        "gpu",
-        "rocm-only",
-    ] + if_google([
-        # Optional: only run internally if ROCm config is enabled
-        "manual",
-    ]),
+    backends = ["gpu"],
+    tags = ["rocm-only"],
     deps = [
         ":rocm_collector",
         ":rocm_tracer",
diff --git a/third_party/xla/xla/codegen/emitters/tests/BUILD b/third_party/xla/xla/codegen/emitters/tests/BUILD
index 7cedec6f546e7c..2a0f84a1a1fdae 100644
--- a/third_party/xla/xla/codegen/emitters/tests/BUILD
+++ b/third_party/xla/xla/codegen/emitters/tests/BUILD
@@ -13,7 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 
-load("@bazel_skylib//rules:copy_file.bzl", "copy_file")
 load("//xla:lit.bzl", "lit_test_suite")
 load(
     "//xla/tsl/platform:build_config_root.bzl",
@@ -26,24 +25,6 @@ package(
     licenses = ["notice"],
 )
 
-copy_file(
-    name = "cpu_fusion_to_mlir_copy",
-    testonly = True,
-    src = "//xla/backends/cpu/codegen/tools:fusion_to_mlir",
-    out = "cpu_fusion_to_mlir",
-    allow_symlink = True,
-    is_executable = True,
-)
-
-copy_file(
-    name = "gpu_fusion_to_mlir_copy",
-    testonly = True,
-    src = "//xla/backends/gpu/codegen/tools:fusion_to_mlir",
-    out = "gpu_fusion_to_mlir",
-    allow_symlink = True,
-    is_executable = True,
-)
-
 #Need no gpu to run
 _FILECHECK_ONLY_TESTS = [
     #TODO(rocm) Most of thease test cuda lowering only
@@ -77,9 +58,9 @@ lit_test_suite(
     default_tags = tf_cuda_tests_tags(),
     exec_properties = tf_exec_properties({"tags": tf_cuda_tests_tags()}),
     tools = [
-        ":cpu_fusion_to_mlir_copy",
-        ":gpu_fusion_to_mlir_copy",
+        "//xla/backends/cpu/codegen/tools:cpu_fusion_to_mlir",
         "//xla/backends/cpu/codegen/tools:cpu_test_correctness",
+        "//xla/backends/gpu/codegen/tools:gpu_fusion_to_mlir",
         "//xla/backends/gpu/codegen/tools:gpu_test_correctness",
         "//xla/codegen/tools:emitters_opt",
         "@llvm-project//llvm:FileCheck",
@@ -92,8 +73,8 @@ lit_test_suite(
     cfg = "//xla:lit.cfg.py",
     default_tags = ["gpu"],
     tools = [
-        ":cpu_fusion_to_mlir_copy",
-        ":gpu_fusion_to_mlir_copy",
+        "//xla/backends/cpu/codegen/tools:cpu_fusion_to_mlir",
+        "//xla/backends/gpu/codegen/tools:gpu_fusion_to_mlir",
         "//xla/codegen/tools:emitters_opt",
         "@llvm-project//llvm:FileCheck",
     ],
diff --git a/third_party/xla/xla/codegen/xtile/codegen/experimental_fusion_emitter.cc b/third_party/xla/xla/codegen/xtile/codegen/experimental_fusion_emitter.cc
index 4e164a0fb08885..bfb664909a058c 100644
--- a/third_party/xla/xla/codegen/xtile/codegen/experimental_fusion_emitter.cc
+++ b/third_party/xla/xla/codegen/xtile/codegen/experimental_fusion_emitter.cc
@@ -1367,6 +1367,24 @@ absl::Status EmitScanComputation(mlir::ImplicitLocOpBuilder& b,
   return absl::OkStatus();
 }
 
+// Returns `init` with the scan dimension re-inserted as a unit dimension, e.g.
+// `tensor<16xf32>` -> `tensor<16x1xf32>` for a scan along dimension 1.
+//
+// The inits and carries of `xtile.scan` keep the scan dimension so that they
+// have the same rank as the scan inputs and outputs. Rank-reduced carries would
+// be rank-zero tensors for 1D scans, which Triton does not support, and the
+// unit dimension lets the lowering broadcast a carry across a tile and extract
+// the next carry from it without any rank changes.
+TensorValue ExpandScanDim(mlir::ImplicitLocOpBuilder& b, TensorValue init,
+                          int64_t scan_dim) {
+  SmallVector<int64_t> shape(init.getType().getShape());
+  shape.insert(shape.begin() + scan_dim, 1);
+  auto type =
+      mlir::RankedTensorType::get(shape, init.getType().getElementType());
+  return mlir::cast<TensorValue>(
+      mlir::stablehlo::ReshapeOp::create(b, type, init).getResult());
+}
+
 absl::StatusOr<std::vector<TensorValue>> EmitScan(
     EmitterContext& emitter_ctx,
     const ge::TiledHloInstruction& tiled_hlo_scan) {
@@ -1445,8 +1463,11 @@ absl::StatusOr<std::vector<TensorValue>> EmitScan(
 
     SmallVector<Value> inits;
     for (int i = 0; i < num_operands; ++i) {
-      inits.push_back(emitter_ctx.TiledHloToTensorValue(
-          *tiled_hlo_scan.operand(num_operands + i)));
+      inits.push_back(
+          ExpandScanDim(b,
+                        emitter_ctx.TiledHloToTensorValue(
+                            *tiled_hlo_scan.operand(num_operands + i)),
+                        scan_dim));
     }
 
     Value zero = MakeIndex(b, 0);
@@ -1456,7 +1477,12 @@ absl::StatusOr<std::vector<TensorValue>> EmitScan(
     b.setInsertionPointToStart(for_op.getBody());
 
     Value iv = for_op.getInductionVar();
-    emitter_ctx.MapSymbolIdToSequentialDimValue(dim_info.id, iv,
+    Value tile_idx = iv;
+    if (hlo_scan.is_reverse()) {
+      Value max_tile = MakeIndex(b, loop_count - 1);
+      tile_idx = arith::SubIOp::create(b, max_tile, iv);
+    }
+    emitter_ctx.MapSymbolIdToSequentialDimValue(dim_info.id, tile_idx,
                                                 Interval{0, loop_count - 1});
 
     const auto& input_region = tiled_hlo_scan.hlo_regions().front();
@@ -1469,7 +1495,7 @@ absl::StatusOr<std::vector<TensorValue>> EmitScan(
     SmallVector<Type> carry_types;
     SmallVector<Type> output_types;
 
-    Value iv_i32 = Cast(b, iv, b.getI32Type());
+    Value tile_idx_i32 = Cast(b, tile_idx, b.getI32Type());
 
     for (int i = 0; i < num_operands; ++i) {
       TensorValue input_tile = input_results[i];
@@ -1483,7 +1509,7 @@ absl::StatusOr<std::vector<TensorValue>> EmitScan(
           CreateConst(b, init_type.getElementType(), 0.0f);
       ABSL_ASSIGN_OR_RETURN(input_tile,
                        MaskOperand(b, *input_region.roots()[i], input_tile,
-                                   iv_i32, scan_dim, neutral_value));
+                                   tile_idx_i32, scan_dim, neutral_value));
       masked_inputs.push_back(input_tile);
     }
 
@@ -1549,8 +1575,11 @@ absl::StatusOr<std::vector<TensorValue>> EmitScan(
         mask_dim_bounds.push_back(input.getType().getDimSize(idx));
       }
     }
-    TensorValue init_tensor = emitter_ctx.TiledHloToTensorValue(
-        *tiled_hlo_scan.operand(num_operands + i));
+    TensorValue init_tensor =
+        ExpandScanDim(b,
+                      emitter_ctx.TiledHloToTensorValue(
+                          *tiled_hlo_scan.operand(num_operands + i)),
+                      scan_dim);
     mlir::Value neutral_value =
         CreateConst(b, init_tensor.getType().getElementType(), 0.0f);
 
diff --git a/third_party/xla/xla/codegen/xtile/ir/tests/ops.mlir b/third_party/xla/xla/codegen/xtile/ir/tests/ops.mlir
index 6087ef26f9838c..4a8457949af424 100644
--- a/third_party/xla/xla/codegen/xtile/ir/tests/ops.mlir
+++ b/third_party/xla/xla/codegen/xtile/ir/tests/ops.mlir
@@ -188,7 +188,19 @@ func.func @illegal_mask_out_of_bounds(%src: tensor<32xf64>, %mask: f64) -> tenso
 
 // -----
 
-func.func @scan_op(%input: tensor<10x20x30xf32>, %init: tensor<10x20xf32>) -> (tensor<10x20x30xf32>, tensor<10x20xf32>) {
+func.func @scan_op(%input: tensor<10x20x30xf32>, %init: tensor<10x20x1xf32>) -> (tensor<10x20x30xf32>, tensor<10x20x1xf32>) {
+  %output, %carry = xtile.scan(%input) inits(%init) dimension = 2 {scan_dim_size = 30 : i64} : (tensor<10x20x30xf32>), (tensor<10x20x1xf32>) -> (tensor<10x20x30xf32>), (tensor<10x20x1xf32>) {
+  ^bb0(%arg0: tensor<10x20xf32>, %arg1: tensor<10x20xf32>):
+    %add = stablehlo.add %arg0, %arg1 : tensor<10x20xf32>
+    stablehlo.return %add : tensor<10x20xf32>
+  }
+  return %output, %carry : tensor<10x20x30xf32>, tensor<10x20x1xf32>
+}
+
+// -----
+
+func.func @scan_op_rank_reduced_init(%input: tensor<10x20x30xf32>, %init: tensor<10x20xf32>) -> (tensor<10x20x30xf32>, tensor<10x20xf32>) {
+  // expected-error @+1 {{init #0 type 'tensor<10x20xf32>' does not match expected type 'tensor<10x20x1xf32>'}}
   %output, %carry = xtile.scan(%input) inits(%init) dimension = 2 {scan_dim_size = 30 : i64} : (tensor<10x20x30xf32>), (tensor<10x20xf32>) -> (tensor<10x20x30xf32>), (tensor<10x20xf32>) {
   ^bb0(%arg0: tensor<10x20xf32>, %arg1: tensor<10x20xf32>):
     %add = stablehlo.add %arg0, %arg1 : tensor<10x20xf32>
@@ -199,6 +211,66 @@ func.func @scan_op(%input: tensor<10x20x30xf32>, %init: tensor<10x20xf32>) -> (t
 
 // -----
 
+func.func @scan_op_dimension_out_of_range(%input: tensor<10x20x30xf32>, %init: tensor<10x20x1xf32>) -> (tensor<10x20x30xf32>, tensor<10x20x1xf32>) {
+  // expected-error @+1 {{scan dimension 3 is out of range for input #0 of rank 3}}
+  %output, %carry = xtile.scan(%input) inits(%init) dimension = 3 {scan_dim_size = 30 : i64} : (tensor<10x20x30xf32>), (tensor<10x20x1xf32>) -> (tensor<10x20x30xf32>), (tensor<10x20x1xf32>) {
+  ^bb0(%arg0: tensor<10x20xf32>, %arg1: tensor<10x20xf32>):
+    %add = stablehlo.add %arg0, %arg1 : tensor<10x20xf32>
+    stablehlo.return %add : tensor<10x20xf32>
+  }
+  return %output, %carry : tensor<10x20x30xf32>, tensor<10x20x1xf32>
+}
+
+// -----
+
+func.func @scan_op_init_count_mismatch(%input: tensor<10x20x30xf32>, %init: tensor<10x20x1xf32>) -> (tensor<10x20x30xf32>, tensor<10x20x1xf32>) {
+  // expected-error @+1 {{expects one init, output and carry per input, but got 1 inputs, 2 inits, 1 outputs and 1 carries}}
+  %output, %carry = xtile.scan(%input) inits(%init, %init) dimension = 2 {scan_dim_size = 30 : i64} : (tensor<10x20x30xf32>), (tensor<10x20x1xf32>, tensor<10x20x1xf32>) -> (tensor<10x20x30xf32>), (tensor<10x20x1xf32>) {
+  ^bb0(%arg0: tensor<10x20xf32>, %arg1: tensor<10x20xf32>):
+    %add = stablehlo.add %arg0, %arg1 : tensor<10x20xf32>
+    stablehlo.return %add : tensor<10x20xf32>
+  }
+  return %output, %carry : tensor<10x20x30xf32>, tensor<10x20x1xf32>
+}
+
+// -----
+
+func.func @scan_op_negative_dimension(%input: tensor<10x20x30xf32>, %init: tensor<10x20x1xf32>) -> (tensor<10x20x30xf32>, tensor<10x20x1xf32>) {
+  // expected-error @+1 {{scan dimension -1 is out of range for input #0 of rank 3}}
+  %output, %carry = xtile.scan(%input) inits(%init) dimension = -1 {scan_dim_size = 30 : i64} : (tensor<10x20x30xf32>), (tensor<10x20x1xf32>) -> (tensor<10x20x30xf32>), (tensor<10x20x1xf32>) {
+  ^bb0(%arg0: tensor<10x20xf32>, %arg1: tensor<10x20xf32>):
+    %add = stablehlo.add %arg0, %arg1 : tensor<10x20xf32>
+    stablehlo.return %add : tensor<10x20xf32>
+  }
+  return %output, %carry : tensor<10x20x30xf32>, tensor<10x20x1xf32>
+}
+
+// -----
+
+func.func @scan_op_output_type_mismatch(%input: tensor<10x20x30xf32>, %init: tensor<10x20x1xf32>) -> (tensor<10x20x30xf64>, tensor<10x20x1xf32>) {
+  // expected-error @+1 {{output #0 type 'tensor<10x20x30xf64>' does not match input type 'tensor<10x20x30xf32>'}}
+  %output, %carry = xtile.scan(%input) inits(%init) dimension = 2 {scan_dim_size = 30 : i64} : (tensor<10x20x30xf32>), (tensor<10x20x1xf32>) -> (tensor<10x20x30xf64>), (tensor<10x20x1xf32>) {
+  ^bb0(%arg0: tensor<10x20xf32>, %arg1: tensor<10x20xf32>):
+    %add = stablehlo.add %arg0, %arg1 : tensor<10x20xf32>
+    stablehlo.return %add : tensor<10x20xf32>
+  }
+  return %output, %carry : tensor<10x20x30xf64>, tensor<10x20x1xf32>
+}
+
+// -----
+
+func.func @scan_op_carry_type_mismatch(%input: tensor<10x20x30xf32>, %init: tensor<10x20x1xf32>) -> (tensor<10x20x30xf32>, tensor<10x20x30xf32>) {
+  // expected-error @+1 {{carry #0 type 'tensor<10x20x30xf32>' does not match expected type 'tensor<10x20x1xf32>'}}
+  %output, %carry = xtile.scan(%input) inits(%init) dimension = 2 {scan_dim_size = 30 : i64} : (tensor<10x20x30xf32>), (tensor<10x20x1xf32>) -> (tensor<10x20x30xf32>), (tensor<10x20x30xf32>) {
+  ^bb0(%arg0: tensor<10x20xf32>, %arg1: tensor<10x20xf32>):
+    %add = stablehlo.add %arg0, %arg1 : tensor<10x20xf32>
+    stablehlo.return %add : tensor<10x20xf32>
+  }
+  return %output, %carry : tensor<10x20x30xf32>, tensor<10x20x30xf32>
+}
+
+// -----
+
 // expected-error @+1 {{layout has 0 dimensions, but shape has 1}}
 func.func @memref_layout_shape_size_mismatch(%arg0: memref<1024xf32, #xtile.layout<[]>>) {
   return
diff --git a/third_party/xla/xla/codegen/xtile/ir/xtile_ops.cc b/third_party/xla/xla/codegen/xtile/ir/xtile_ops.cc
index 2151e7895ac053..aada99430a4ce4 100644
--- a/third_party/xla/xla/codegen/xtile/ir/xtile_ops.cc
+++ b/third_party/xla/xla/codegen/xtile/ir/xtile_ops.cc
@@ -24,6 +24,7 @@ limitations under the License.
 #include "absl/log/check.h"
 #include "absl/strings/string_view.h"
 #include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Builders.h"
@@ -241,6 +242,54 @@ mlir::OpFoldResult MaskOp::fold(FoldAdaptor) {
   return {};
 }
 
+mlir::LogicalResult ScanOp::verify() {
+  mlir::ValueRange inputs = getInputs();
+  if (getInits().size() != inputs.size() ||
+      getOutputs().size() != inputs.size() ||
+      getCarries().size() != inputs.size()) {
+    return emitOpError() << "expects one init, output and carry per input, but "
+                            "got "
+                         << inputs.size() << " inputs, " << getInits().size()
+                         << " inits, " << getOutputs().size() << " outputs and "
+                         << getCarries().size() << " carries";
+  }
+
+  int64_t dimension = getDimension();
+  for (auto [index, input] : llvm::enumerate(inputs)) {
+    mlir::RankedTensorType input_type =
+        mlir::cast<mlir::RankedTensorType>(input.getType());
+    if (dimension < 0 || dimension >= input_type.getRank()) {
+      return emitOpError() << "scan dimension " << dimension
+                           << " is out of range for input #" << index
+                           << " of rank " << input_type.getRank();
+    }
+    if (getOutputs()[index].getType() != input_type) {
+      return emitOpError() << "output #" << index << " type "
+                           << getOutputs()[index].getType()
+                           << " does not match input type " << input_type;
+    }
+
+    // The inits and carries keep the scan dimension as a unit dimension, so
+    // that they have the same rank as the inputs and outputs.
+    llvm::SmallVector<int64_t> carry_shape(input_type.getShape());
+    carry_shape[dimension] = 1;
+    mlir::RankedTensorType carry_type =
+        mlir::RankedTensorType::get(carry_shape, input_type.getElementType());
+    if (getInits()[index].getType() != carry_type) {
+      return emitOpError() << "init #" << index << " type "
+                           << getInits()[index].getType()
+                           << " does not match expected type " << carry_type;
+    }
+    if (getCarries()[index].getType() != carry_type) {
+      return emitOpError() << "carry #" << index << " type "
+                           << getCarries()[index].getType()
+                           << " does not match expected type " << carry_type;
+    }
+  }
+
+  return mlir::success();
+}
+
 static bool IsSupportedScaleElementType(mlir::Type type) {
   return mlir::isa<mlir::Float8E8M0FNUType, mlir::Float8E4M3FNType,
                    mlir::Float8E5M2Type>(type) ||
diff --git a/third_party/xla/xla/codegen/xtile/ir/xtile_ops.td b/third_party/xla/xla/codegen/xtile/ir/xtile_ops.td
index cc3c19ab9fcad0..aa03b558989d4a 100644
--- a/third_party/xla/xla/codegen/xtile/ir/xtile_ops.td
+++ b/third_party/xla/xla/codegen/xtile/ir/xtile_ops.td
@@ -296,6 +296,14 @@ def ScanOp : XTile_Op<"scan", [Pure, AttrSizedOperandSegments, AttrSizedResultSe
   let summary = "scan operation";
   let description = [{
     Mirrors HLO's kScan operation.
+
+    Unlike HLO, the inits and carries keep the scan dimension as a unit
+    dimension, i.e. they have the same rank as the inputs and outputs. This
+    avoids rank-zero carries for 1D scans, which Triton does not support, and
+    allows broadcasting a carry across a tile without changing its rank.
+
+    Example: for inputs of type `tensor<16x1024xf32>` and `dimension = 1`, the
+    inits and carries have type `tensor<16x1xf32>`.
   }];
 
   let arguments = (
@@ -321,6 +329,8 @@ def ScanOp : XTile_Op<"scan", [Pure, AttrSizedOperandSegments, AttrSizedResultSe
     `:` `(` type($inputs) `)` `,` `(` type($inits) `)` `->` `(` type($outputs) `)` `,` `(` type($carries) `)`
     $body
   }];
+
+  let hasVerifier = 1;
 }
 
 //
diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
index d5bc83cbac0e8c..b2737e63898297 100644
--- a/third_party/xla/xla/service/gpu/BUILD
+++ b/third_party/xla/xla/service/gpu/BUILD
@@ -3546,7 +3546,6 @@ cc_library(
         ":launch_dimensions",
         "//xla:future",
         "//xla:status_macros",
-        "//xla:util",
         "//xla/codegen/emitters:computation_fingerprint",
         "//xla/codegen/emitters:kernel_arguments",
         "//xla/hlo/ir:hlo",
@@ -3554,6 +3553,8 @@ cc_library(
         "//xla/stream_executor/gpu:tma_metadata",
         "//xla/tsl/concurrency:future",
         "//xla/tsl/platform:env",
+        "//xla/tsl/platform:logging",
+        "//xla/tsl/util:sorted_range",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
@@ -3578,9 +3579,7 @@ xla_cc_test(
         ":kernel_reuse_cache_proto_cc",
         "//xla/tests:xla_internal_test_main",
         "//xla/tsl/concurrency:future",
-        "//xla/tsl/lib/core:status_test_util",
         "//xla/tsl/platform:env",
-        "//xla/tsl/platform:statusor",
         "//xla/tsl/util/proto:proto_matchers",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
diff --git a/third_party/xla/xla/service/gpu/gpu_compiler.cc b/third_party/xla/xla/service/gpu/gpu_compiler.cc
index 94e4cb7b6b7bb2..69c73b3f4c3f43 100644
--- a/third_party/xla/xla/service/gpu/gpu_compiler.cc
+++ b/third_party/xla/xla/service/gpu/gpu_compiler.cc
@@ -2034,15 +2034,16 @@ void AddGemmRewriterPasses(HloPassPipeline& pipeline,
     bias_mode = GemmRewriterOptions::BiasMode::kNoBias;
   }
 
+  GemmRewriterOptions fp8_options{GemmRewriterOptions::DType::kFp8Only,
+                                  bias_mode};
+  pipeline.AddPass<GemmRewriter>(gpu_version, toolkit_version, fp8_options);
+
   // Rewrite dots with the algorithms that cannot be handled by cublas directly.
   // I.e. transform single dot into a chain of dots with the default algorithm
   // that cublas can handle. These dots were inlined by the CallInliner pass
   // above.
   pipeline.AddPass<DotAlgorithmRewriter>();
 
-  GemmRewriterOptions fp8_options{GemmRewriterOptions::DType::kFp8Only,
-                                  bias_mode};
-  pipeline.AddPass<GemmRewriter>(gpu_version, toolkit_version, fp8_options);
   pipeline.AddPass<GemmRewriter>(
       gpu_version, toolkit_version,
       GemmRewriterOptions{GemmRewriterOptions::DType::kNonFp8Only, bias_mode});
diff --git a/third_party/xla/xla/service/gpu/kernel_reuse_cache.cc b/third_party/xla/xla/service/gpu/kernel_reuse_cache.cc
index e99faacf0c73c8..06437a34b9e234 100644
--- a/third_party/xla/xla/service/gpu/kernel_reuse_cache.cc
+++ b/third_party/xla/xla/service/gpu/kernel_reuse_cache.cc
@@ -46,7 +46,8 @@ limitations under the License.
 #include "xla/tsl/concurrency/future.h"
 #include "xla/tsl/platform/env.h"
 #include "xla/tsl/platform/file_system.h"
-#include "xla/util.h"
+#include "xla/tsl/platform/logging.h"
+#include "xla/tsl/util/sorted_range.h"
 
 namespace xla::gpu {
 namespace {
@@ -88,15 +89,17 @@ absl::Status KernelReuseCache::Load(const CompilationCacheProto& proto) {
     return absl::OkStatus();
   }
   absl::MutexLock lock(m_);
-  for (const auto& [name, entry] : proto.entries()) {
+  for (const auto& [name, entry] : tsl::KeySortedRange(proto.entries())) {
     std::optional<se::ClusterDim> cluster_dim;
     if (entry.has_cluster_dim()) {
       cluster_dim =
           se::ClusterDim{entry.cluster_dim().x(), entry.cluster_dim().y(),
                          entry.cluster_dim().z()};
     }
-    std::vector<uint8_t> binary(entry.binary().data(),
-                                entry.binary().data() + entry.binary().size());
+    std::shared_ptr<const std::vector<uint8_t>> binary =
+        std::make_shared<const std::vector<uint8_t>>(
+            entry.binary().data(),
+            entry.binary().data() + entry.binary().size());
     TF_RET_CHECK(
         cache_
             .insert(
@@ -116,7 +119,7 @@ CompilationCacheProto KernelReuseCache::Export() const {
   absl::MutexLock lock(m_);
   CompilationCacheProto proto;
   proto.set_compatibility_version(kCacheCompatibilityVersion);
-  for (const auto& [fingerprint, future] : cache_) {
+  for (const auto& [fingerprint, future] : tsl::KeySortedRange(cache_)) {
     const absl::StatusOr<Entry>& cache_entry = future.Await();
     if (!cache_entry.ok()) {
       // If a generator failed, the Future will hold an error.
@@ -147,9 +150,11 @@ CompilationCacheProto KernelReuseCache::Export() const {
       *proto_entry.mutable_cluster_dim() = cluster_dim_proto;
     }
     proto_entry.set_shmem_bytes(cache_entry->shmem_bytes);
-    proto_entry.set_binary(absl::string_view(
-        reinterpret_cast<const char*>(cache_entry->binary.data()),
-        cache_entry->binary.size()));
+    if (cache_entry->binary != nullptr) {
+      proto_entry.set_binary(absl::string_view(
+          reinterpret_cast<const char*>(cache_entry->binary->data()),
+          cache_entry->binary->size()));
+    }
   }
   return proto;
 }
@@ -168,12 +173,13 @@ absl::Status UpdateDiskKernelCache(absl::string_view path, const bool do_append,
   }
 
   absl::flat_hash_set<std::string> kernel_fingerprints;
-  for (const auto& [_, entry] : disk_cache.entries()) {
+  for (const auto& [_, entry] : tsl::KeySortedRange(disk_cache.entries())) {
     kernel_fingerprints.insert(entry.fingerprint());
   }
 
   int stored_kernel_count = 0;
-  for (const auto& [name, entry] : current_cache.entries()) {
+  for (const auto& [name, entry] :
+       tsl::KeySortedRange(current_cache.entries())) {
     if (kernel_fingerprints.contains(entry.fingerprint())) {
       continue;
     }
@@ -190,7 +196,7 @@ absl::Status UpdateDiskKernelCache(absl::string_view path, const bool do_append,
   return absl::OkStatus();
 }
 
-std::pair<tsl::Future<const KernelReuseCache::Entry*>, bool>
+std::pair<tsl::Future<KernelReuseCache::Entry>, bool>
 KernelReuseCache::GetWithStatus(
     const HloComputation* fused_computation,
     absl::Span<const emitters::KernelArgument> kernel_arguments,
@@ -203,7 +209,7 @@ KernelReuseCache::GetWithStatus(
   return GetWithStatus(std::move(fingerprint), generator);
 }
 
-std::pair<tsl::Future<const KernelReuseCache::Entry*>, bool>
+std::pair<tsl::Future<KernelReuseCache::Entry>, bool>
 KernelReuseCache::GetWithStatus(
     std::string fingerprint,
     absl::FunctionRef<tsl::Future<KernelReuseCache::Entry>()> generator) {
@@ -219,9 +225,7 @@ KernelReuseCache::GetWithStatus(
     it = cache_.insert({std::move(fingerprint), generator()}).first;
   }
 
-  return {it->second.Map(
-              [](const KernelReuseCache::Entry& entry) { return &entry; }),
-          cached};
+  return {it->second, cached};
 }
 
 }  // namespace xla::gpu
diff --git a/third_party/xla/xla/service/gpu/kernel_reuse_cache.h b/third_party/xla/xla/service/gpu/kernel_reuse_cache.h
index d887d461058fef..0f99f8c8b1ce5c 100644
--- a/third_party/xla/xla/service/gpu/kernel_reuse_cache.h
+++ b/third_party/xla/xla/service/gpu/kernel_reuse_cache.h
@@ -16,6 +16,7 @@ limitations under the License.
 #define XLA_SERVICE_GPU_KERNEL_REUSE_CACHE_H_
 
 #include <cstdint>
+#include <memory>
 #include <optional>
 #include <string>
 #include <utility>
@@ -49,7 +50,11 @@ class KernelReuseCache {
     LaunchDimensions launch_dimensions;
     std::optional<se::ClusterDim> cluster_dim;
     int64_t shmem_bytes = 0;
-    std::vector<uint8_t> binary;
+    // Reference counted so that consumers can share the binary instead of
+    // copying it. Default-initialized to an empty vector so `binary` is
+    // non-null even for default-constructed entries.
+    std::shared_ptr<const std::vector<uint8_t>> binary =
+        std::make_shared<const std::vector<uint8_t>>();
     stream_executor::gpu::TmaMetadata tma_metadata;
     bool use_pdl = false;
   };
@@ -77,9 +82,7 @@ class KernelReuseCache {
   // Retrieves the cache entry for the given computation, or generates it
   // asynchronously using the given generator function and stores it in the
   // cache.
-  //
-  // The returned pointer is never nullptr.
-  std::pair<tsl::Future<const Entry*>, bool /*was_cached*/> GetWithStatus(
+  std::pair<tsl::Future<Entry>, bool /*was_cached*/> GetWithStatus(
       const HloComputation* fused_computation,
       absl::Span<const emitters::KernelArgument> kernel_arguments,
       absl::string_view discriminator,
@@ -89,11 +92,9 @@ class KernelReuseCache {
   // asynchronously using the given generator function and stores it in the
   // cache.
   //
-  // The returned pointer is never nullptr.
-  //
   // A non-OK status is returned if the entry is not found and the generator
   // failed.
-  std::pair<tsl::Future<const Entry*>, bool /*was_cached*/> GetWithStatus(
+  std::pair<tsl::Future<Entry>, bool /*was_cached*/> GetWithStatus(
       std::string fingerprint,
       absl::FunctionRef<tsl::Future<Entry>()> generator);
 
diff --git a/third_party/xla/xla/service/gpu/kernel_reuse_cache_test.cc b/third_party/xla/xla/service/gpu/kernel_reuse_cache_test.cc
index f2749399f69b3a..79e349f48638f2 100644
--- a/third_party/xla/xla/service/gpu/kernel_reuse_cache_test.cc
+++ b/third_party/xla/xla/service/gpu/kernel_reuse_cache_test.cc
@@ -14,7 +14,12 @@ limitations under the License.
 ==============================================================================*/
 #include "xla/service/gpu/kernel_reuse_cache.h"
 
+#include <cstdint>
+#include <memory>
 #include <string>
+#include <tuple>
+#include <utility>
+#include <vector>
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
@@ -22,9 +27,7 @@ limitations under the License.
 #include "absl/status/status.h"
 #include "xla/service/gpu/kernel_reuse_cache.pb.h"
 #include "xla/tsl/concurrency/future.h"
-#include "xla/tsl/lib/core/status_test_util.h"
 #include "xla/tsl/platform/env.h"
-#include "xla/tsl/platform/statusor.h"
 #include "xla/tsl/util/proto/proto_matchers.h"
 
 namespace xla::gpu {
@@ -47,9 +50,8 @@ TEST_F(KernelReuseTest, ExportAndLoadWork) {
   EXPECT_FALSE(future.IsReady());
   promise.Set(KernelReuseCache::Entry{kernel_name});
 
-  ASSERT_OK_AND_ASSIGN(const KernelReuseCache::Entry* result, future.Await());
-  EXPECT_THAT(result, testing::NotNull());
-  EXPECT_EQ(result->kernel_name, kernel_name);
+  ASSERT_OK_AND_ASSIGN(const KernelReuseCache::Entry& result, future.Await());
+  EXPECT_EQ(result.kernel_name, kernel_name);
   EXPECT_FALSE(cache.IsEmpty());
 
   const CompilationCacheProto proto = cache.Export();
@@ -76,9 +78,8 @@ TEST_F(KernelReuseTest, ExportAndLoadWork) {
       return absl::UnimplementedError("Should be cached");
     });
     EXPECT_TRUE(was_cached);
-    ASSERT_OK_AND_ASSIGN(const KernelReuseCache::Entry* result, future.Await());
-    EXPECT_THAT(result, testing::NotNull());
-    EXPECT_EQ(result->kernel_name, kernel_name);
+    ASSERT_OK_AND_ASSIGN(const KernelReuseCache::Entry& result, future.Await());
+    EXPECT_EQ(result.kernel_name, kernel_name);
   }
 }
 
@@ -89,8 +90,11 @@ TEST_F(KernelReuseTest, UpdatingDiskKernelCacheWorks) {
     const CompilationCacheProto proto = [](std::string kernel_name) {
       KernelReuseCache cache;
       auto [result, was_cached] = cache.GetWithStatus("fingerprint", [&]() {
-        return KernelReuseCache::Entry{.kernel_name = kernel_name,
-                                       .binary = {5, 6}};
+        KernelReuseCache::Entry entry;
+        entry.kernel_name = kernel_name;
+        entry.binary = std::make_shared<const std::vector<uint8_t>>(
+            std::vector<uint8_t>{5, 6});
+        return entry;
       });
       return cache.Export();
     }("k1");
@@ -101,8 +105,11 @@ TEST_F(KernelReuseTest, UpdatingDiskKernelCacheWorks) {
     const CompilationCacheProto proto = [](std::string kernel_name) {
       KernelReuseCache cache;
       auto [result, was_cached] = cache.GetWithStatus("fingerprint1", [&]() {
-        return KernelReuseCache::Entry{.kernel_name = kernel_name,
-                                       .binary = {7, 8}};
+        KernelReuseCache::Entry entry;
+        entry.kernel_name = kernel_name;
+        entry.binary = std::make_shared<const std::vector<uint8_t>>(
+            std::vector<uint8_t>{7, 8});
+        return entry;
       });
       return cache.Export();
     }("k2");
@@ -117,5 +124,85 @@ TEST_F(KernelReuseTest, UpdatingDiskKernelCacheWorks) {
   EXPECT_EQ(proto.entries_size(), 2);
 }
 
+TEST_F(KernelReuseTest, EntryRemainsValidAfterCacheClearAndDestruction) {
+  auto [promise, returned] = tsl::MakePromise<KernelReuseCache::Entry>();
+  tsl::Future<KernelReuseCache::Entry> future;
+  tsl::Future<KernelReuseCache::Entry> cached_future;
+  {
+    KernelReuseCache cache;
+    bool was_cached = false;
+    std::tie(future, was_cached) = cache.GetWithStatus(
+        "fingerprint1", [returned = std::move(returned)]() mutable {
+          return std::move(returned);
+        });
+    EXPECT_FALSE(was_cached);
+
+    std::tie(cached_future, was_cached) = cache.GetWithStatus(
+        "fingerprint1",
+        []() { return absl::UnimplementedError("Should be cached"); });
+    EXPECT_TRUE(was_cached);
+
+    cache.Clear();
+  }
+  {
+    auto p = std::move(promise);
+    KernelReuseCache::Entry expected;
+    expected.kernel_name = "kernel_after_clear";
+    expected.binary = std::make_shared<const std::vector<uint8_t>>(
+        std::vector<uint8_t>{1, 2, 3});
+    p.Set(std::move(expected));
+  }
+
+  ASSERT_OK_AND_ASSIGN(const KernelReuseCache::Entry& result, future.Await());
+  EXPECT_EQ(result.kernel_name, "kernel_after_clear");
+  ASSERT_NE(result.binary, nullptr);
+  EXPECT_THAT(*result.binary, testing::ElementsAre(1, 2, 3));
+
+  ASSERT_OK_AND_ASSIGN(const KernelReuseCache::Entry& cached_result,
+                       cached_future.Await());
+  EXPECT_EQ(cached_result.binary.get(), result.binary.get());
+}
+
+TEST_F(KernelReuseTest, DefaultConstructedEntryHasNonNullEmptyBinary) {
+  KernelReuseCache::Entry entry{"empty_kernel"};
+  ASSERT_NE(entry.binary, nullptr);
+  EXPECT_TRUE(entry.binary->empty());
+}
+
+TEST_F(KernelReuseTest, MovingFromReturnedEntryDoesNotMutateCachedEntry) {
+  KernelReuseCache cache;
+  auto [first_future, first_cached] =
+      cache.GetWithStatus("fp", []() -> tsl::Future<KernelReuseCache::Entry> {
+        KernelReuseCache::Entry entry;
+        entry.kernel_name = "shared_kernel";
+        entry.binary = std::make_shared<const std::vector<uint8_t>>(
+            std::vector<uint8_t>{10, 20, 30});
+        return entry;
+      });
+  EXPECT_FALSE(first_cached);
+
+  // Simulate a consumer (like TritonFusion::Emit) copying Entry into a result
+  // struct and then std::move'ing kernel_name and binary out of its copy.
+  auto moved_future = first_future.Map([](KernelReuseCache::Entry entry_copy) {
+    std::string moved_name = std::move(entry_copy.kernel_name);
+    std::shared_ptr<const std::vector<uint8_t>> moved_binary =
+        std::move(entry_copy.binary);
+    return std::make_pair(std::move(moved_name), std::move(moved_binary));
+  });
+  ASSERT_OK_AND_ASSIGN(auto first_moved, moved_future.Await());
+  EXPECT_EQ(first_moved.first, "shared_kernel");
+  ASSERT_NE(first_moved.second, nullptr);
+
+  // A subsequent cache hit must still see the intact kernel_name and share the
+  // exact same underlying binary buffer.
+  auto [second_future, second_cached] = cache.GetWithStatus(
+      "fp", []() { return absl::UnimplementedError("Should be cached"); });
+  EXPECT_TRUE(second_cached);
+  ASSERT_OK_AND_ASSIGN(const KernelReuseCache::Entry& second_result,
+                       second_future.Await());
+  EXPECT_EQ(second_result.kernel_name, "shared_kernel");
+  EXPECT_EQ(second_result.binary.get(), first_moved.second.get());
+}
+
 }  // namespace
 }  // namespace xla::gpu
diff --git a/third_party/xla/xla/service/gpu/thunk_emitter.cc b/third_party/xla/xla/service/gpu/thunk_emitter.cc
index 4b855978761bc1..b17bfa3db9b318 100644
--- a/third_party/xla/xla/service/gpu/thunk_emitter.cc
+++ b/third_party/xla/xla/service/gpu/thunk_emitter.cc
@@ -1601,13 +1601,14 @@ Future<ThunkSequence> ThunkEmitter::EmitTritonCustomCall(
                     tma_metadata = result.tma_metadata,
                     kernel_name = std::move(kernel_name)](
                        const std::vector<uint8_t>& cubin) mutable {
-                return KernelReuseCache::Entry{std::move(kernel_name),
-                                               launch_dimensions,
-                                               /*cluster_dim=*/std::nullopt,
-                                               shmem_bytes,
-                                               cubin,
-                                               tma_metadata,
-                                               use_pdl};
+                return KernelReuseCache::Entry{
+                    std::move(kernel_name),
+                    launch_dimensions,
+                    /*cluster_dim=*/std::nullopt,
+                    shmem_bytes,
+                    std::make_shared<const std::vector<uint8_t>>(cubin),
+                    tma_metadata,
+                    use_pdl};
               });
         });
   };
@@ -1626,19 +1627,19 @@ Future<ThunkSequence> ThunkEmitter::EmitTritonCustomCall(
   return status_or_entry.Map(
       [info = std::move(info), kernel_arguments = std::move(kernel_arguments),
        call_zeroed_outputs = std::move(call_zeroed_outputs)](
-          const KernelReuseCache::Entry* entry) mutable
+          const KernelReuseCache::Entry& entry) mutable
           -> absl::StatusOr<ThunkSequence> {
-        ABSL_ASSIGN_OR_RETURN(CustomKernel custom_kernel,
-                         kernel::CreateOwnedCubinCustomKernel(
-                             entry->kernel_name, entry->binary,
-                             kernel_arguments.args().size(),
-                             entry->launch_dimensions.block_counts(),
-                             entry->launch_dimensions.thread_counts_per_block(),
-                             entry->shmem_bytes));
+        ABSL_ASSIGN_OR_RETURN(
+            CustomKernel custom_kernel,
+            kernel::CreateSharedCubinCustomKernel(
+                entry.kernel_name, entry.binary, kernel_arguments.args().size(),
+                entry.launch_dimensions.block_counts(),
+                entry.launch_dimensions.thread_counts_per_block(),
+                entry.shmem_bytes));
         return ThunkSequence::Of<CustomKernelThunk>(
             std::move(info), std::move(custom_kernel),
-            std::move(kernel_arguments), entry->use_pdl, call_zeroed_outputs,
-            entry->tma_metadata);
+            std::move(kernel_arguments), entry.use_pdl, call_zeroed_outputs,
+            entry.tma_metadata);
       });
 }
 
diff --git a/third_party/xla/xla/service/spmd/custom_call_handler.cc b/third_party/xla/xla/service/spmd/custom_call_handler.cc
index 89557f9b38f4ab..f6d8472633c488 100644
--- a/third_party/xla/xla/service/spmd/custom_call_handler.cc
+++ b/third_party/xla/xla/service/spmd/custom_call_handler.cc
@@ -142,9 +142,12 @@ constexpr char kSPMDOpMultiSlice[] = "_SPMDInternalOp_MultiSlice";
 constexpr char kSPMDOpWrap[] = "_SPMDInternalOp_Wrap";
 
 bool IsTopKIndexOutputUsed(HloInstruction* hlo) {
+  if (hlo->IsRoot()) {
+    return true;
+  }
   for (const HloInstruction* user : hlo->users()) {
     if (user->opcode() == HloOpcode::kGetTupleElement) {
-      if (user->tuple_index() > 0 && !user->users().empty()) {
+      if (user->tuple_index() > 0 && !user->IsDead()) {
         return true;
       }
     } else {
diff --git a/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc b/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc
index 5e1ca3c5e817cd..e5b3c990657169 100644
--- a/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc
+++ b/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc
@@ -17450,6 +17450,55 @@ ENTRY %entry {
   EXPECT_THAT(topk_operand, op::Shape("bf16[32,256000]{1,0}"));
 }
 
+TEST_P(SpmdPartitioningTest, TopKCustomCallTupleRootUsesIndices) {
+  absl::string_view hlo_string = R"(
+HloModule module
+
+region {
+  index_lhs = s32[] parameter(2)
+  index_rhs = s32[] parameter(3)
+  value_lhs = bf16[] parameter(0)
+  value_rhs = bf16[] parameter(1)
+  ROOT compare = pred[] compare(value_lhs, value_rhs), direction=GT, order=TOTAL
+}
+
+ENTRY entry {
+  input = bf16[64,4096]{1,0} parameter(0), sharding={devices=[2,1]<=[2]}
+  ROOT topk = (bf16[64,40]{1,0}, s32[64,40]{1,0}) custom-call(input), custom_call_target="TopK", called_computations={region}, sharding={{devices=[2,1]<=[2]}, {devices=[2,1]<=[2]}}
+})";
+
+  ASSERT_OK_AND_ASSIGN(auto module,
+                       PartitionComputation(hlo_string, /*num_devices=*/2));
+  VLOG(1) << module->ToString();
+  EXPECT_THAT(module->entry_computation()->root_instruction(),
+              op::Shape("(bf16[32,40]{1,0}, s32[32,40]{1,0})"));
+}
+
+TEST_P(SpmdPartitioningTest, TopKCustomCallIndexRootUsesIndices) {
+  absl::string_view hlo_string = R"(
+HloModule module
+
+region {
+  index_lhs = s32[] parameter(2)
+  index_rhs = s32[] parameter(3)
+  value_lhs = bf16[] parameter(0)
+  value_rhs = bf16[] parameter(1)
+  ROOT compare = pred[] compare(value_lhs, value_rhs), direction=GT, order=TOTAL
+}
+
+ENTRY entry {
+  input = bf16[64,4096]{1,0} parameter(0), sharding={devices=[2,1]<=[2]}
+  topk = (bf16[64,40]{1,0}, s32[64,40]{1,0}) custom-call(input), custom_call_target="TopK", called_computations={region}, sharding={{devices=[2,1]<=[2]}, {devices=[2,1]<=[2]}}
+  ROOT indices = s32[64,40]{1,0} get-tuple-element(topk), index=1, sharding={devices=[2,1]<=[2]}
+})";
+
+  ASSERT_OK_AND_ASSIGN(auto module,
+                       PartitionComputation(hlo_string, /*num_devices=*/2));
+  VLOG(1) << module->ToString();
+  EXPECT_THAT(module->entry_computation()->root_instruction(),
+              op::Shape("s32[32,40]{1,0}"));
+}
+
 TEST_P(SpmdPartitioningTest,
        TopKCustomCallTopkReplicatedOperandNonTopKDimSharded) {
   absl::string_view hlo_string = R"(
diff --git a/third_party/xla/xla/stream_executor/BUILD b/third_party/xla/xla/stream_executor/BUILD
index 15ba54aaf48e28..fbfaaba6f928b0 100644
--- a/third_party/xla/xla/stream_executor/BUILD
+++ b/third_party/xla/xla/stream_executor/BUILD
@@ -881,7 +881,6 @@ cc_library(
         ":kernel_args",
         ":kernel_args_packing_spec",
         ":kernel_spec_proto_cc",
-        "//xla/tsl/platform:statusor",
         "@com_google_absl//absl/functional:function_ref",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
diff --git a/third_party/xla/xla/stream_executor/kernel_spec.cc b/third_party/xla/xla/stream_executor/kernel_spec.cc
index b689bb7f171d7d..d2ea62ce1f7d5d 100644
--- a/third_party/xla/xla/stream_executor/kernel_spec.cc
+++ b/third_party/xla/xla/stream_executor/kernel_spec.cc
@@ -17,6 +17,7 @@ limitations under the License.
 
 #include <cstddef>
 #include <cstdint>
+#include <memory>
 #include <optional>
 #include <string>
 #include <utility>
@@ -31,7 +32,6 @@ limitations under the License.
 #include "absl/types/span.h"
 #include "xla/stream_executor/kernel_args_packing_spec.h"
 #include "xla/stream_executor/kernel_spec.pb.h"
-#include "xla/tsl/platform/statusor.h"
 
 namespace stream_executor {
 
@@ -64,6 +64,15 @@ KernelLoaderSpec KernelLoaderSpec::CreateOwningCudaCubinInMemorySpec(
                           std::move(kernel_name), arity, kernel_args_packing};
 }
 
+KernelLoaderSpec KernelLoaderSpec::CreateSharedCudaCubinInMemorySpec(
+    std::shared_ptr<const std::vector<uint8_t>> cubin_bytes,
+    std::string kernel_name, size_t arity,
+    KernelArgsPacking kernel_args_packing) {
+  CHECK(cubin_bytes != nullptr);
+  return KernelLoaderSpec{SharedCudaCubinInMemory{std::move(cubin_bytes)},
+                          std::move(kernel_name), arity, kernel_args_packing};
+}
+
 KernelLoaderSpec KernelLoaderSpec::CreateCudaPtxInMemorySpec(
     absl::string_view ptx, std::string kernel_name, size_t arity,
     KernelArgsPacking kernel_args_packing) {
@@ -78,6 +87,14 @@ KernelLoaderSpec KernelLoaderSpec::CreateOwningCudaPtxInMemorySpec(
                           std::move(kernel_name), arity, kernel_args_packing};
 }
 
+KernelLoaderSpec KernelLoaderSpec::CreateSharedCudaPtxInMemorySpec(
+    std::shared_ptr<const std::string> ptx, std::string kernel_name,
+    size_t arity, KernelArgsPacking kernel_args_packing) {
+  CHECK(ptx != nullptr);
+  return KernelLoaderSpec{SharedCudaPtxInMemory{std::move(ptx)},
+                          std::move(kernel_name), arity, kernel_args_packing};
+}
+
 absl::StatusOr<KernelLoaderSpecProto> KernelLoaderSpec::ToProto() const {
   if (std::holds_alternative<KernelArgsPackingFunc>(kernel_args_packing_) &&
       std::get<KernelArgsPackingFunc>(kernel_args_packing_) != nullptr) {
diff --git a/third_party/xla/xla/stream_executor/kernel_spec.h b/third_party/xla/xla/stream_executor/kernel_spec.h
index f25a790cc55b9b..69eee23a8ffbe9 100644
--- a/third_party/xla/xla/stream_executor/kernel_spec.h
+++ b/third_party/xla/xla/stream_executor/kernel_spec.h
@@ -77,6 +77,15 @@ struct OwningCudaPtxInMemory {
   std::string ptx;
 };
 
+// Like OwningCudaPtxInMemory, but the PTX is held in a reference counted buffer
+// that can be shared between multiple loader specs. Prefer this over
+// OwningCudaPtxInMemory when the same kernel is referenced from several places
+// (e.g. an HLO module that invokes the same kernel multiple times), so that the
+// PTX is stored only once.
+struct SharedCudaPtxInMemory {
+  std::shared_ptr<const std::string> ptx;
+};
+
 // Kernel loader specification for a CUBIN blob that resides in memory.
 struct CudaCubinInMemory {
   absl::Span<const uint8_t> cubin_bytes;
@@ -87,6 +96,15 @@ struct OwningCudaCubinInMemory {
   std::vector<uint8_t> cubin_bytes;
 };
 
+// Like OwningCudaCubinInMemory, but the CUBIN data is held in a reference
+// counted buffer that can be shared between multiple loader specs. Prefer this
+// over OwningCudaCubinInMemory when the same kernel is referenced from several
+// places (e.g. an HLO module that invokes the same kernel multiple times), so
+// that the CUBIN is stored only once.
+struct SharedCudaCubinInMemory {
+  std::shared_ptr<const std::vector<uint8_t>> cubin_bytes;
+};
+
 // Describes how to load a kernel on any subset of a number of target platforms.
 class KernelLoaderSpec {
  public:
@@ -114,11 +132,13 @@ class KernelLoaderSpec {
   }
   bool has_cuda_cubin_in_memory() const {
     return std::holds_alternative<CudaCubinInMemory>(payload_) ||
-           std::holds_alternative<OwningCudaCubinInMemory>(payload_);
+           std::holds_alternative<OwningCudaCubinInMemory>(payload_) ||
+           std::holds_alternative<SharedCudaCubinInMemory>(payload_);
   }
   bool has_cuda_ptx_in_memory() const {
     return std::holds_alternative<CudaPtxInMemory>(payload_) ||
-           std::holds_alternative<OwningCudaPtxInMemory>(payload_);
+           std::holds_alternative<OwningCudaPtxInMemory>(payload_) ||
+           std::holds_alternative<SharedCudaPtxInMemory>(payload_);
   }
 
   // Accessors for platform variant kernel load specifications.
@@ -137,6 +157,10 @@ class KernelLoaderSpec {
       return CudaCubinInMemory{
           std::get<OwningCudaCubinInMemory>(payload_).cubin_bytes};
     }
+    if (std::holds_alternative<SharedCudaCubinInMemory>(payload_)) {
+      return CudaCubinInMemory{
+          *std::get<SharedCudaCubinInMemory>(payload_).cubin_bytes};
+    }
     return std::nullopt;
   }
 
@@ -147,6 +171,9 @@ class KernelLoaderSpec {
     if (std::holds_alternative<OwningCudaPtxInMemory>(payload_)) {
       return CudaPtxInMemory{std::get<OwningCudaPtxInMemory>(payload_).ptx};
     }
+    if (std::holds_alternative<SharedCudaPtxInMemory>(payload_)) {
+      return CudaPtxInMemory{*std::get<SharedCudaPtxInMemory>(payload_).ptx};
+    }
     return std::nullopt;
   }
 
@@ -168,12 +195,24 @@ class KernelLoaderSpec {
   static KernelLoaderSpec CreateOwningCudaCubinInMemorySpec(
       std::vector<uint8_t> cubin_bytes, std::string kernel_name, size_t arity,
       KernelArgsPacking kernel_args_packing = KernelArgsPackingFunc{});
+  // Like CreateOwningCudaCubinInMemorySpec, but the CUBIN buffer can be shared
+  // with other specs. `cubin_bytes` must not be null.
+  static KernelLoaderSpec CreateSharedCudaCubinInMemorySpec(
+      std::shared_ptr<const std::vector<uint8_t>> cubin_bytes,
+      std::string kernel_name, size_t arity,
+      KernelArgsPacking kernel_args_packing = KernelArgsPackingFunc{});
   static KernelLoaderSpec CreateCudaPtxInMemorySpec(
       absl::string_view ptx, std::string kernel_name, size_t arity,
       KernelArgsPacking kernel_args_packing = KernelArgsPackingFunc{});
   static KernelLoaderSpec CreateOwningCudaPtxInMemorySpec(
       std::string ptx, std::string kernel_name, size_t arity,
       KernelArgsPacking kernel_args_packing = KernelArgsPackingFunc{});
+  // Like CreateOwningCudaPtxInMemorySpec, but the PTX buffer can be shared with
+  // other specs. `ptx` must not be null.
+  static KernelLoaderSpec CreateSharedCudaPtxInMemorySpec(
+      std::shared_ptr<const std::string> ptx, std::string kernel_name,
+      size_t arity,
+      KernelArgsPacking kernel_args_packing = KernelArgsPackingFunc{});
 
   void set_kernel_args_packing(KernelArgsPacking kernel_args_packing) {
     kernel_args_packing_ = std::move(kernel_args_packing);
@@ -197,7 +236,8 @@ class KernelLoaderSpec {
  private:
   using Payload =
       std::variant<InProcessSymbol, CudaCubinInMemory, CudaPtxInMemory,
-                   OwningCudaCubinInMemory, OwningCudaPtxInMemory>;
+                   OwningCudaCubinInMemory, OwningCudaPtxInMemory,
+                   SharedCudaCubinInMemory, SharedCudaPtxInMemory>;
 
   explicit KernelLoaderSpec(
       Payload payload, std::string kernel_name, size_t arity,
diff --git a/third_party/xla/xla/stream_executor/kernel_spec_test.cc b/third_party/xla/xla/stream_executor/kernel_spec_test.cc
index 25753c1384465e..97a88ff64fa285 100644
--- a/third_party/xla/xla/stream_executor/kernel_spec_test.cc
+++ b/third_party/xla/xla/stream_executor/kernel_spec_test.cc
@@ -17,6 +17,7 @@ limitations under the License.
 
 #include <array>
 #include <cstdint>
+#include <memory>
 #include <string>
 #include <utility>
 #include <vector>
@@ -118,6 +119,49 @@ TEST(KernelLoaderSpec, OwningCudaPtx) {
   EXPECT_THAT(spec.kernel_name(), "kernel24");
 }
 
+TEST(KernelLoaderSpec, SharedCudaCubin) {
+  static constexpr std::array<uint8_t, 4> kCubinData = {0xDE, 0xAD, 0xBE, 0xEF};
+  auto shared_cubin = std::make_shared<const std::vector<uint8_t>>(
+      kCubinData.begin(), kCubinData.end());
+  auto spec = KernelLoaderSpec::CreateSharedCudaCubinInMemorySpec(
+      shared_cubin, "kernel24", /*arity=*/2);
+  KernelLoaderSpec copy = spec;
+
+  EXPECT_TRUE(spec.has_cuda_cubin_in_memory());
+  EXPECT_FALSE(spec.has_cuda_ptx_in_memory());
+  EXPECT_FALSE(spec.has_in_process_symbol());
+
+  ASSERT_TRUE(spec.cuda_cubin_in_memory().has_value());
+  ASSERT_TRUE(copy.cuda_cubin_in_memory().has_value());
+  EXPECT_THAT(spec.cuda_cubin_in_memory(),
+              Optional(Field(&CudaCubinInMemory::cubin_bytes, kCubinData)));
+  EXPECT_EQ(spec.cuda_cubin_in_memory()->cubin_bytes.data(),
+            shared_cubin->data());
+  EXPECT_EQ(copy.cuda_cubin_in_memory()->cubin_bytes.data(),
+            shared_cubin->data());
+  EXPECT_THAT(spec.kernel_name(), "kernel24");
+}
+
+TEST(KernelLoaderSpec, SharedCudaPtx) {
+  static constexpr absl::string_view kPtxData = "PTX DEADBEEF";
+  auto shared_ptx = std::make_shared<const std::string>(kPtxData);
+  auto spec = KernelLoaderSpec::CreateSharedCudaPtxInMemorySpec(
+      shared_ptx, "kernel24", /*arity=*/2);
+  KernelLoaderSpec copy = spec;
+
+  EXPECT_FALSE(spec.has_cuda_cubin_in_memory());
+  EXPECT_TRUE(spec.has_cuda_ptx_in_memory());
+  EXPECT_FALSE(spec.has_in_process_symbol());
+
+  ASSERT_TRUE(spec.cuda_ptx_in_memory().has_value());
+  ASSERT_TRUE(copy.cuda_ptx_in_memory().has_value());
+  EXPECT_THAT(spec.cuda_ptx_in_memory(),
+              Optional(Field(&CudaPtxInMemory::ptx, kPtxData)));
+  EXPECT_EQ(spec.cuda_ptx_in_memory()->ptx.data(), shared_ptx->data());
+  EXPECT_EQ(copy.cuda_ptx_in_memory()->ptx.data(), shared_ptx->data());
+  EXPECT_THAT(spec.kernel_name(), "kernel24");
+}
+
 TEST(KernelLoaderSpec, PtxKernelFromProto) {
   auto proto = ParseTextProtoOrDie<KernelLoaderSpecProto>(R"pb(
     ptx { data: "PTX!" }
@@ -172,6 +216,30 @@ TEST(KernelLoaderSpec, CubinKernelToProto) {
               )pb")));
 }
 
+TEST(KernelLoaderSpec, SharedPtxKernelToProto) {
+  auto spec = KernelLoaderSpec::CreateSharedCudaPtxInMemorySpec(
+      std::make_shared<const std::string>("PTX!"), "kernel_name", /*arity=*/42);
+
+  EXPECT_THAT(spec.ToProto(), IsOkAndHolds(EqualsProto(R"pb(
+                ptx { data: "PTX!" }
+                kernel_name: "kernel_name"
+                arity: 42
+              )pb")));
+}
+
+TEST(KernelLoaderSpec, SharedCubinKernelToProto) {
+  auto spec = KernelLoaderSpec::CreateSharedCudaCubinInMemorySpec(
+      std::make_shared<const std::vector<uint8_t>>(
+          std::vector<uint8_t>{'C', 'U', 'B', 'I', 'N'}),
+      "kernel_name", /*arity=*/42);
+
+  EXPECT_THAT(spec.ToProto(), IsOkAndHolds(EqualsProto(R"pb(
+                cubin { data: "CUBIN" }
+                kernel_name: "kernel_name"
+                arity: 42
+              )pb")));
+}
+
 TEST(KernelLoaderSpec, InProcessSymbolFromProto) {
   auto proto = ParseTextProtoOrDie<KernelLoaderSpecProto>(R"pb(
     in_process_symbol { persistent_name: "persistent_kernel_name" }
diff --git a/third_party/xla/xla/tests/aot_compatibility_experimental/BUILD b/third_party/xla/xla/tests/aot_compatibility_experimental/BUILD
index 7e6056da533a33..f023698959ff34 100644
--- a/third_party/xla/xla/tests/aot_compatibility_experimental/BUILD
+++ b/third_party/xla/xla/tests/aot_compatibility_experimental/BUILD
@@ -32,7 +32,10 @@ xla_cc_test(
     name = "test_lib_test",
     srcs = ["test_lib_test.cc"],
     data = ["//xla/tests/aot_compatibility_experimental/gpu:test_dummy_test_data"],
-    tags = ["no_oss"],
+    tags = [
+        "broken",
+        "no_oss",
+    ],
     deps = [
         ":test_lib",
         "//xla/pjrt/proto:compile_options_proto_cc",
diff --git a/third_party/xla/xla/tsl/platform/default/build_config_root.bzl b/third_party/xla/xla/tsl/platform/default/build_config_root.bzl
index f31fb4b211de98..bd70e6cef200b2 100644
--- a/third_party/xla/xla/tsl/platform/default/build_config_root.bzl
+++ b/third_party/xla/xla/tsl/platform/default/build_config_root.bzl
@@ -77,6 +77,9 @@ def tf_cuda_2gpu_tests_tags():
 def tf_has_tag(kwargs, tag):
     return ("tags" in kwargs and kwargs["tags"] != None and tag in kwargs["tags"])
 
+def tf_has_any_tag(kwargs, tags):
+    return any([tf_has_tag(kwargs, tag) for tag in tags])
+
 def tf_exec_properties(kwargs):
     """Gets execution_properties for TensorFlow GPU tests based on the provided tags.
 
@@ -85,7 +88,7 @@ def tf_exec_properties(kwargs):
     Returns:
         execution_properties with the execution pool names for rbe.
     """
-    if is_rocm_configured():
+    if is_rocm_configured() and tf_has_any_tag(kwargs, ["requires-gpu-rocm", "requires-gpu-amd"]):
         if tf_has_tag(kwargs, "multi_gpu"):
             return ROCM_MULTI_GPU_TEST_PROPERTIES
         if tf_has_tag(kwargs, "gpu"):