diff --git a/configure.py b/configure.py
index 2bb9c197acf9f2..25d70c1bfe56f7 100644
--- a/configure.py
+++ b/configure.py
@@ -913,8 +913,16 @@ def retrieve_clang_version(clang_executable):
                            stderr=stderr)
 
   curr_version_split = curr_version.lower().split('clang version ')
-  if len(curr_version_split) > 1:
-    curr_version = curr_version_split[1].split()[0].split('git')
+  if len(curr_version_split) <= 1:
+    sys.stdout.write('WARNING: current clang installation version unknown.\n')
+    return None
+
+  tokens = curr_version_split[1].split()
+  if not tokens:
+    sys.stdout.write('WARNING: current clang installation version unknown.\n')
+    return None
+
+  curr_version = tokens[0].split('git')
 
   if len(curr_version) > 1:
     print('WARNING: current clang installation is not a release version.\n')
@@ -937,7 +945,13 @@ def retrieve_clang_version(clang_executable):
 # offset of in the current version of ubp. See
 # https://github.com/protocolbuffers/upb/blob/9effcbcb27f0a665f9f345030188c0b291e32482/upb/upb.c#L183.
 def disable_clang_offsetof_extension(clang_version):
-  if int(clang_version.split('.')[0]) in (16, 17):
+  if not clang_version:
+    return
+  try:
+    clang_major_version = int(clang_version.split('.')[0])
+  except ValueError:
+    return
+  if clang_major_version in (16, 17):
     write_to_bazelrc('build --copt=-Wno-gnu-offsetof-extensions')
 
 
diff --git a/tensorflow/compiler/mlir/lite/flatbuffer_export.cc b/tensorflow/compiler/mlir/lite/flatbuffer_export.cc
index 7d34e71a38aa15..ddf93115835684 100644
--- a/tensorflow/compiler/mlir/lite/flatbuffer_export.cc
+++ b/tensorflow/compiler/mlir/lite/flatbuffer_export.cc
@@ -4518,6 +4518,7 @@ absl::Status Translator::AppendBufferData() {
   for (const auto& [index, buffer] : const_buffer_storage_.buffers()) {
     uint64_t hash = buffer->hash();
     if (hashcode_to_pos.find(hash) == hashcode_to_pos.end()) {
+      export_stream_.get().write_zeros(kFbAlignment - offset() % kFbAlignment);
       int64_t size = 0;
       int64_t buffer_offset = offset();
       auto status = buffer->ApplyData([this, &size](absl::string_view data) {
diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/cluster_ops_by_policy.h b/tensorflow/compiler/mlir/tensorflow/transforms/cluster_ops_by_policy.h
index e3c0ee5cc23238..1c25d9c00bc55a 100644
--- a/tensorflow/compiler/mlir/tensorflow/transforms/cluster_ops_by_policy.h
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/cluster_ops_by_policy.h
@@ -200,7 +200,7 @@ class ClusteringPolicySet {
  private:
   template <typename T, typename... Args>
   void AddImpl(Args&&... args) {
-    static_assert(std::is_base_of<ClusteringPolicy, T>::value,
+    static_assert(std::is_base_of_v<ClusteringPolicy, T>,
                   "T must implement ClusteringPolicy");
     policies_.emplace_back(std::make_unique<T>(std::forward<Args>(args)...));
   }
diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/fused_kernel_matcher.cc b/tensorflow/compiler/mlir/tensorflow/transforms/fused_kernel_matcher.cc
index b2ab71fa5129cb..88be1de8c034ad 100644
--- a/tensorflow/compiler/mlir/tensorflow/transforms/fused_kernel_matcher.cc
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/fused_kernel_matcher.cc
@@ -199,7 +199,7 @@ class FuseContractionWithBiasAdd : public OpRewritePattern<SrcOpT> {
     attrs.push_back(
         NamedAttribute(StringAttr::get(context, "epsilon"), epsilon));
 
-    if (std::is_same<FusedOpT, _FusedConv2DOp>::value) {
+    if (std::is_same_v<FusedOpT, _FusedConv2DOp>) {
       // Here TArgs types do not include types of the first two parameters,
       // i.e. the convolution input and the filter. TArgs are parameters for
       // the extras like the bias etc.
diff --git a/tensorflow/core/data/service/dispatcher_impl.cc b/tensorflow/core/data/service/dispatcher_impl.cc
index cb480c8832a774..5a79e479d00097 100644
--- a/tensorflow/core/data/service/dispatcher_impl.cc
+++ b/tensorflow/core/data/service/dispatcher_impl.cc
@@ -205,12 +205,16 @@ absl::Status ValidateDatasetId(const std::string& dataset_id) {
         absl::StrCat("Invalid dataset ID: ", dataset_id,
                      ". Dataset IDs must not contain '/'."));
   }
+  if (absl::StrContains(dataset_id, '\\')) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("Invalid dataset ID: ", dataset_id,
+                     ". Dataset IDs must not contain '\\'."));
+  }
 #if defined(_WIN32)
-  if (absl::StrContains(dataset_id, '\\') ||
-      absl::StrContains(dataset_id, ':')) {
+  if (absl::StrContains(dataset_id, ':')) {
     return absl::InvalidArgumentError(
         absl::StrCat("Invalid dataset ID: ", dataset_id,
-                     ". Dataset IDs must not contain '\\' or ':'."));
+                     ". Dataset IDs must not contain ':'."));
   }
 #endif
   return absl::OkStatus();
diff --git a/tensorflow/core/data/service/grpc_dispatcher_impl_test.cc b/tensorflow/core/data/service/grpc_dispatcher_impl_test.cc
index e68dc402565c55..85eff83dcda965 100644
--- a/tensorflow/core/data/service/grpc_dispatcher_impl_test.cc
+++ b/tensorflow/core/data/service/grpc_dispatcher_impl_test.cc
@@ -175,6 +175,24 @@ TEST_F(GrpcDispatcherImplTest, GetSplitInvalidProviderIndex) {
   }
 }
 
+TEST_F(GrpcDispatcherImplTest, GetOrRegisterDatasetInvalidDatasetId) {
+  const std::vector<std::string> invalid_ids = {
+      "..\\..\\etc\\passwd", "a\\b", "a/b", ".", "..",
+  };
+
+  for (const auto& id : invalid_ids) {
+    ClientContext ctx;
+    GetOrRegisterDatasetRequest req;
+    *req.mutable_dataset()->mutable_graph() = testing::RangeDataset(10).graph();
+    req.set_dataset_id(id);
+    GetOrRegisterDatasetResponse resp;
+    ::grpc::Status status =
+        dispatcher_client_stub_->GetOrRegisterDataset(&ctx, req, &resp);
+    EXPECT_EQ(status.error_code(), ::grpc::StatusCode::INVALID_ARGUMENT)
+        << "Dataset ID '" << id << "' should be rejected.";
+  }
+}
+
 }  // namespace
 }  // namespace data
 }  // namespace tensorflow
diff --git a/tensorflow/core/grappler/optimizers/constant_folding.cc b/tensorflow/core/grappler/optimizers/constant_folding.cc
index 6801d61262ae42..e0be95530d5ab1 100644
--- a/tensorflow/core/grappler/optimizers/constant_folding.cc
+++ b/tensorflow/core/grappler/optimizers/constant_folding.cc
@@ -2042,17 +2042,6 @@ void ConstantFolding::ReplaceBinaryOperationWithBroadcastTo(
   graph_modified_ = true;
 }
 
-void ConstantFolding::ReplaceDivisionOfOnesByReciprocal(NodeDef* node,
-                                                        GraphDef* graph) {
-  node->set_op("Reciprocal");
-  node->mutable_input()->SwapElements(0, 1);
-  const std::string ctrl_dep =
-      AddControlDependency(node->input(1), graph, node_map_.get());
-  node_map_->UpdateInput(node->name(), node->input(1), ctrl_dep);
-  node->set_input(1, ctrl_dep);
-  graph_modified_ = true;
-}
-
 void ConstantFolding::ReplaceSubtractionFromZeroByNegation(NodeDef* node,
                                                            GraphDef* graph) {
   node->set_op("Neg");
@@ -3047,15 +3036,10 @@ absl::Status ConstantFolding::SimplifyArithmeticOperations(
       return absl::OkStatus();
     }
 
-    // Replace 1 / y with Reciprocal op.
-    if (y_matches_output_shape && is_any_div && x_is_one) {
-      TF_RETURN_IF_ERROR(CheckAttrExists(*node, "T"));
-      DataType type = node->attr().at("T").type();
-      if (DataTypeIsFloating(type) || DataTypeIsComplex(type)) {
-        ReplaceDivisionOfOnesByReciprocal(node, optimized_graph);
-        return absl::OkStatus();
-      }
-    }
+    // Note: 1 / y is intentionally not rewritten to Reciprocal(y). The CPU
+    // Reciprocal kernel uses Eigen's fast-math reciprocal for float, which
+    // is not exactly IEEE division, so the rewrite silently changed results
+    // between eager and graph execution on x86 (see issue #102771).
 
     const bool y_is_zero = IsZeros(*y);
     const bool y_is_one = y_is_zero ? false : IsOnes(*y);
diff --git a/tensorflow/core/grappler/optimizers/constant_folding.h b/tensorflow/core/grappler/optimizers/constant_folding.h
index 2ac90cd681a1b8..5e9113414aa408 100644
--- a/tensorflow/core/grappler/optimizers/constant_folding.h
+++ b/tensorflow/core/grappler/optimizers/constant_folding.h
@@ -131,7 +131,6 @@ class ConstantFolding : public GraphOptimizer {
                                                   NodeDef* node,
                                                   GraphDef* graph);
 
-  void ReplaceDivisionOfOnesByReciprocal(NodeDef* node, GraphDef* graph);
   absl::Status FoldGraph(
       const GraphProperties& properties, GraphDef* output,
       absl::flat_hash_set<std::string>* nodes_to_not_simplify);
diff --git a/tensorflow/core/grappler/optimizers/constant_folding_test.cc b/tensorflow/core/grappler/optimizers/constant_folding_test.cc
index 49137d15140bcc..8ef638accac8ab 100644
--- a/tensorflow/core/grappler/optimizers/constant_folding_test.cc
+++ b/tensorflow/core/grappler/optimizers/constant_folding_test.cc
@@ -888,9 +888,11 @@ TEST_F(ConstantFoldingTest, NeutralElement) {
         EXPECT_EQ("x", node.input(0));
         EXPECT_EQ(ctrl_ones_name, node.input(1));
       } else if (name == "div2") {
-        EXPECT_EQ("Reciprocal", node.op());
-        EXPECT_EQ("y", node.input(0));
-        EXPECT_EQ(ctrl_ones_name, node.input(1));
+        // ones / y is not rewritten to Reciprocal(y): the CPU Reciprocal
+        // kernel is not exactly IEEE division for float (see issue #102771).
+        EXPECT_EQ("Div", node.op());
+        EXPECT_EQ(ones_name, node.input(0));
+        EXPECT_EQ("y", node.input(1));
       } else if (name == "floordiv") {
         EXPECT_EQ("FloorDiv", node.op());
         EXPECT_EQ("x", node.input(0));
diff --git a/tensorflow/core/util/tensor_slice_reader_cache.cc b/tensorflow/core/util/tensor_slice_reader_cache.cc
index ddb3e36d1e6dbe..1eb3410ae581c1 100644
--- a/tensorflow/core/util/tensor_slice_reader_cache.cc
+++ b/tensorflow/core/util/tensor_slice_reader_cache.cc
@@ -52,8 +52,6 @@ TensorSliceReaderCache::~TensorSliceReaderCache() {
 const TensorSliceReader* TensorSliceReaderCache::GetReader(
     const std::string& filepattern,
     TensorSliceReader::OpenTableFunction open_function, int preferred_shard) {
-  mutex_lock l(mu_);
-
 #if defined(__GXX_RTTI) || defined(_CPPRTTI)
   // Get the function pointer from the open_function value.
   TensorSliceReaderCache::OpenFuncType* func_ptr =
@@ -72,22 +70,43 @@ const TensorSliceReader* TensorSliceReaderCache::GetReader(
     return nullptr;
   }
 
-  // Wait if another thread is already trying to open the same files.
-  while (still_opening_.find(filepattern) != still_opening_.end()) {
-    cv_.wait(l);
+  TensorSliceReader* reader = nullptr;
+
+  // scope block for lock
+  {
+    mutex_lock l(mu_);
+
+    // Wait if another thread is already trying to open the same files.
+    while (still_opening_.find(filepattern) != still_opening_.end()) {
+      cv_.wait(l);
+    }
+
+    auto it = readers_.find(filepattern);
+    if (it != readers_.end()) {
+      auto cached_val = it->second;
+      if (cached_val.first == *func_ptr) {
+        reader = cached_val.second;
+        VLOG(1) << "Using cached TensorSliceReader for " << filepattern << ": "
+                << reader;
+      } else {
+        LOG(WARNING) << "Caching disabled because the checkpoint file "
+                     << "is being opened with two different open functions: "
+                     << filepattern;
+      }
+      return reader;
+    }
+
+    still_opening_.insert(filepattern);
   }
 
-  TensorSliceReader* reader = nullptr;
-  if (readers_.find(filepattern) == readers_.end()) {
+  // no lock for expensive constructing TensorSliceReader
+  TensorSliceReader* tmp_reader(
+      new TensorSliceReader(filepattern, open_function, preferred_shard));
+
+  // scope block for lock
+  {
     VLOG(1) << "Creating new TensorSliceReader for " << filepattern;
-    still_opening_.insert(filepattern);
-    // Release the lock temporary as constructing TensorSliceReader is
-    // expensive.
-    mu_.unlock();
-    TensorSliceReader* tmp_reader(
-        new TensorSliceReader(filepattern, open_function, preferred_shard));
-    // Acquire the lock again.
-    mu_.lock();
+    mutex_lock l(mu_);
     if (tmp_reader->status().ok()) {
       reader = tmp_reader;
       readers_[filepattern] = std::make_pair(*func_ptr, reader);
@@ -96,20 +115,9 @@ const TensorSliceReader* TensorSliceReaderCache::GetReader(
     }
     CHECK_EQ(size_t{1}, still_opening_.erase(filepattern));
     VLOG(1) << "Cached TensorSliceReader for " << filepattern << ": " << reader;
-  } else {
-    auto cached_val = readers_[filepattern];
-    if (cached_val.first == *func_ptr) {
-      reader = cached_val.second;
-      VLOG(1) << "Using cached TensorSliceReader for " << filepattern << ": "
-              << reader;
-    } else {
-      LOG(WARNING) << "Caching disabled because the checkpoint file "
-                   << "is being opened with two different open functions: "
-                   << filepattern;
-    }
-  }
 
-  cv_.notify_all();
+    cv_.notify_all();
+  }
   return reader;
 }
 
diff --git a/tensorflow/lite/python/lite_v2_test.py b/tensorflow/lite/python/lite_v2_test.py
index ae73b2046f6731..1afeb4ca0ea32b 100644
--- a/tensorflow/lite/python/lite_v2_test.py
+++ b/tensorflow/lite/python/lite_v2_test.py
@@ -5592,6 +5592,86 @@ def testCOncreteFunctionFloat(self):
     actual_value = self._evaluateTFLiteModel(tflite_model, [input_data])
     self.assertEqual(expected_value.numpy(), actual_value)
 
+  @test_util.run_v2_only
+  def testUseBufferOffsetAlignment(self):
+    """Test that all external constant buffers are aligned to 16 bytes."""
+
+    class MultiConstantModel(tf.Module):
+
+      def __init__(self):
+        super().__init__()
+        # Constants with non-multiple-of-16 byte sizes to test padding alignment
+        self.w1 = tf.Variable(
+            tf.ones([3, 1], dtype=tf.float32), name='w1'
+        )  # 3 * 4 = 12 bytes
+        self.w2 = tf.Variable(
+            tf.ones([5, 1], dtype=tf.float32), name='w2'
+        )  # 5 * 4 = 20 bytes
+        self.w3 = tf.Variable(
+            tf.ones([1, 1], dtype=tf.float32), name='w3'
+        )  # 1 * 4 = 4 bytes
+        self.w4 = tf.Variable(
+            tf.ones([7, 1], dtype=tf.float32), name='w4'
+        )  # 7 * 4 = 28 bytes
+        self.w5 = tf.Variable(
+            tf.ones([6, 1], dtype=tf.float32), name='w5'
+        )  # 6 * 4 = 24 bytes
+
+      @tf.function
+      def __call__(self, x1, x2, x3, x4, x5):
+        return (
+            tf.matmul(x1, self.w1)
+            + tf.matmul(x2, self.w2)
+            + tf.matmul(x3, self.w3)
+            + tf.matmul(x4, self.w4)
+            + tf.matmul(x5, self.w5)
+        )
+
+    root = MultiConstantModel()
+    inputs = [
+        tf.constant([[1.0, 2.0, 3.0]], dtype=tf.float32),
+        tf.constant([[1.0, 2.0, 3.0, 4.0, 5.0]], dtype=tf.float32),
+        tf.constant([[1.0]], dtype=tf.float32),
+        tf.constant([[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0]], dtype=tf.float32),
+        tf.constant([[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]], dtype=tf.float32),
+    ]
+    concrete_func = root.__call__.get_concrete_function(*inputs)
+
+    converter = lite.TFLiteConverterV2.from_concrete_functions(
+        [concrete_func], root
+    )
+    converter._experimental_use_buffer_offset = True
+    tflite_model = converter.convert()
+
+    # Parse flatbuffer model and check all buffer offsets
+    model_obj = schema_fb.Model.GetRootAsModel(tflite_model, 0)
+    external_buffer_count = 0
+    for i in range(model_obj.BuffersLength()):
+      buf = model_obj.Buffers(i)
+      if buf.Offset() > 1:
+        external_buffer_count += 1
+        self.assertEqual(
+            buf.Offset() % 16,
+            0,
+            f'Buffer {i} offset {buf.Offset()} is not 16-byte aligned (offset %'
+            f' 16 = {buf.Offset() % 16})',
+        )
+    self.assertGreaterEqual(external_buffer_count, 5)
+
+    # Evaluate converted model
+    expected_value = root(*inputs)
+    interp = interpreter.Interpreter(model_content=tflite_model)
+    runner = interp.get_signature_runner()
+    output = runner(
+        x1=inputs[0],
+        x2=inputs[1],
+        x3=inputs[2],
+        x4=inputs[3],
+        x5=inputs[4],
+    )
+    actual_value = list(output.values())[0]
+    self.assertEqual(expected_value.numpy(), actual_value)
+
   @test_util.run_v2_only
   def testConcreteFunctionStringInput(self):
     class Model(tf.Module):
diff --git a/tensorflow/python/_pywrap_tensorflow.def b/tensorflow/python/_pywrap_tensorflow.def
index 70aca4798f4b53..937fd5cc377552 100644
--- a/tensorflow/python/_pywrap_tensorflow.def
+++ b/tensorflow/python/_pywrap_tensorflow.def
@@ -256,6 +256,7 @@ EXPORTS
  ?ComputeGradient@Tape@gradients@tensorflow@@QEAA?AVStatus@lts_20260526@absl@@PEAVAbstractContext@3@V?$Span@QEAVAbstractTensorHandle@tensorflow@@@56@11V?$Span@PEAVAbstractTensorHandle@tensorflow@@@56@@Z
  ?Convert@PythonTensorConverter@tensorflow@@QEBA?AV?$unique_ptr@U_object@@UPyDecrefDeleter@detail@tensorflow@@@std@@PEAU_object@@AEAW4DataType@2@PEA_N@Z
  ?Convert@tflite@@YAPEAU_object@@PEAU2@00_N0PEBVPyFunctionLibrary@quantization@tensorflow@@@Z
+ ?ConvertMlirBytecode@tflite@@YAPEAU_object@@PEAU2@00@Z
  ?ConvertPyObjectToAttributeType@tensorflow@@YA?AV?$unique_ptr@U_object@@UPyDecrefDeleter@detail@tensorflow@@@std@@PEAU_object@@W4AttributeType@1@@Z
  ?ConvertPythonAPIParameters@tensorflow@@YA_NAEBVPythonAPIInfo@1@AEBVPythonTensorConverter@1@V?$Span@PEAU_object@@@lts_20260526@absl@@PEAUInferredAttributes@21@@Z
  ?ConvertToEagerTensor@tensorflow@@YAPEAUTFE_TensorHandle@@PEAUTFE_Context@@PEAU_object@@W4DataType@1@PEBD@Z
diff --git a/tensorflow/python/data/kernel_tests/BUILD b/tensorflow/python/data/kernel_tests/BUILD
index 0ae595a8f34275..12291416374d26 100644
--- a/tensorflow/python/data/kernel_tests/BUILD
+++ b/tensorflow/python/data/kernel_tests/BUILD
@@ -1422,6 +1422,7 @@ py_test(
         ":checkpoint_test_base",
         ":test_base",
         "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:debug_mode",
         "//tensorflow/python/eager:def_function",
         "//tensorflow/python/framework:combinations",
         "//tensorflow/python/ops:variables",
diff --git a/tensorflow/python/data/kernel_tests/io_test.py b/tensorflow/python/data/kernel_tests/io_test.py
index ad2945dce322d5..e888149435d41c 100644
--- a/tensorflow/python/data/kernel_tests/io_test.py
+++ b/tensorflow/python/data/kernel_tests/io_test.py
@@ -21,9 +21,11 @@
 
 from absl.testing import parameterized
 import numpy as np
+
 from tensorflow.python.data.kernel_tests import checkpoint_test_base
 from tensorflow.python.data.kernel_tests import test_base
 from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import debug_mode
 from tensorflow.python.eager import def_function
 from tensorflow.python.framework import combinations
 from tensorflow.python.ops import variables
@@ -70,6 +72,15 @@ def testCardinality(self):
     dataset2 = dataset_ops.Dataset.load(self._test_dir, dataset.element_spec)
     self.assertEqual(self.evaluate(dataset2.cardinality()), 42)
 
+  @combinations.generate(test_base.eager_only_combinations())
+  def testSaveInDebugModeWithoutShardFunction(self):
+    debug_mode.toggle_debug_mode(True)
+    self.addCleanup(debug_mode.toggle_debug_mode, False)
+    dataset = dataset_ops.Dataset.range(42)
+    self.evaluate(dataset.save(self._test_dir))
+    dataset2 = dataset_ops.Dataset.load(self._test_dir, dataset.element_spec)
+    self.assertDatasetProduces(dataset2, range(42))
+
   @combinations.generate(test_base.default_test_combinations())
   def testCustomShardFunction(self):
     dataset = dataset_ops.Dataset.range(42)
diff --git a/tensorflow/python/data/ops/save_op.py b/tensorflow/python/data/ops/save_op.py
index c5a63477aeeb3f..5cd78a93b502b4 100644
--- a/tensorflow/python/data/ops/save_op.py
+++ b/tensorflow/python/data/ops/save_op.py
@@ -97,7 +97,7 @@ def set_save_dataset_attributes(dataset, shard_func, path):
   """Sets parameters for SaveDatasetOp and SaveDatasetV2Op."""
   if shard_func is None:
     use_shard_func = False
-    shard_func = lambda *x: None  # a dummy function that will not be used
+    shard_func = lambda *x: 0  # a dummy function that will not be used
   else:
     use_shard_func = True
   wrapped_func = structured_function.StructuredFunctionWrapper(
diff --git a/tensorflow/python/grappler/constant_folding_test.py b/tensorflow/python/grappler/constant_folding_test.py
index 9a0e9fd73b73e1..71f940fc1360c8 100644
--- a/tensorflow/python/grappler/constant_folding_test.py
+++ b/tensorflow/python/grappler/constant_folding_test.py
@@ -102,6 +102,25 @@ def f(x, y):
     self.assertEqual(assign_count, 1)
     self.assertLen(graphs[0].node, 11)
 
+  # See GitHub issue #102771.
+  def testDivisionOfOnesNotRewrittenToReciprocal(self):
+    # 1 / x must stay a true division in the optimized graph. The CPU
+    # Reciprocal kernel uses Eigen's fast-math reciprocal for float, which
+    # is not exactly IEEE division, so rewriting 1 / x to Reciprocal(x)
+    # changed results between eager and graph execution on x86.
+
+    @def_function.function
+    def div_by_ones(x):
+      return 1.0 / x
+
+    with context.eager_mode():
+      ones = array_ops.ones([64], dtype=dtypes.float32)
+      with context.collect_graphs(optimized=True) as graphs:
+        result = div_by_ones(ones).numpy()
+    self.assertLen(graphs, 1)
+    self.assertNotIn('Reciprocal', [node.op for node in graphs[0].node])
+    self.assertAllEqual(result, np.ones([64], dtype=np.float32))
+
 
 if __name__ == '__main__':
   test.main()
diff --git a/tensorflow/python/ops/control_flow_ops_test.py b/tensorflow/python/ops/control_flow_ops_test.py
index 7e536c49648aa1..5fb2757295cd70 100644
--- a/tensorflow/python/ops/control_flow_ops_test.py
+++ b/tensorflow/python/ops/control_flow_ops_test.py
@@ -1661,6 +1661,25 @@ def testWhileLoopSameReturnShape_TrueSingleLoopVar(self):
         c, b, [i], return_same_structure=True, maximum_iterations=50)
     self.assertEqual(self.evaluate(r), [10])
 
+  @test_util.run_v2_only
+  def testEagerWhileLoopSingleLoopVarBareTensorBodyPreservesShape(self):
+    x = constant_op.constant([[5.0]], shape=[1, 1])
+    shapes_seen = []
+
+    def cond(x):
+      shapes_seen.append(x.shape.as_list())
+      return math_ops.greater(x[0, 0], 3)
+
+    def body(x):
+      return x - 1
+
+    r = while_loop.while_loop(cond, body, [x], return_same_structure=True)
+
+    self.assertAllEqual(shapes_seen, [[1, 1]] * len(shapes_seen))
+    self.assertLen(shapes_seen, 3)
+    self.assertIsInstance(r, list)
+    self.assertAllClose(self.evaluate(r), [[[3.0]]])
+
   @test_util.enable_control_flow_v2
   @test_util.run_in_graph_and_eager_modes
   def testSkipsUnnecessaryCaptureGradients(self):
diff --git a/tensorflow/python/ops/while_loop.py b/tensorflow/python/ops/while_loop.py
index d1964dcfdf437a..690ebd3974c9b0 100644
--- a/tensorflow/python/ops/while_loop.py
+++ b/tensorflow/python/ops/while_loop.py
@@ -482,13 +482,18 @@ def while_loop(cond,
     if executing_eagerly:
       packed = False  # whether the body result was packed into a 1-item tuple
 
+      orig_loop_vars_type = (
+          type(loop_vars) if type(loop_vars) in (list, tuple) else list
+      )
+
       loop_var_structure = nest.map_structure(type_spec.type_spec_from_value,
                                               list(loop_vars))
       while cond(*loop_vars):
         loop_vars = body(*loop_vars)
-        if try_to_pack and not isinstance(loop_vars, (list, tuple)):
-          packed = True
-          loop_vars = (loop_vars,)
+        if not isinstance(loop_vars, (list, tuple)):
+          if try_to_pack:
+            packed = True
+          loop_vars = orig_loop_vars_type((loop_vars,))
         nest.assert_same_structure(loop_var_structure, list(loop_vars))
 
       def convert(x):
diff --git a/tensorflow/tools/def_file_filter/symbols_pybind.txt b/tensorflow/tools/def_file_filter/symbols_pybind.txt
index 003a77aa2965b9..f8fce5e7322775 100644
--- a/tensorflow/tools/def_file_filter/symbols_pybind.txt
+++ b/tensorflow/tools/def_file_filter/symbols_pybind.txt
@@ -259,6 +259,7 @@ tensorflow::ImportFunction
 
 [//tensorflow/compiler/mlir/lite/python:converter_python_api] # converter_python_api
 tflite::Convert
+tflite::ConvertMlirBytecode
 tflite::MlirQuantizeModel
 tflite::MlirSparsifyModel
 tflite::RegisterCustomOpdefs
diff --git a/tensorflow/workspace0.bzl b/tensorflow/workspace0.bzl
index dec1ab868435ee..890e8896f0a97f 100644
--- a/tensorflow/workspace0.bzl
+++ b/tensorflow/workspace0.bzl
@@ -40,16 +40,12 @@ def workspace():
     models_repositories()
     bazel_toolchains_repositories()
 
-    # Apple rules for Bazel. https://github.com/bazelbuild/rules_apple.
-    # Note: We add this to fix Kokoro builds.
-    # The rules below call into `rules_proto` but the hash has changed and
-    # Bazel refuses to continue. So, we add our own mirror.
     tf_http_archive(
         name = "rules_proto",
-        sha256 = "20b240eba17a36be4b0b22635aca63053913d5c1ee36e16be36499d167a2f533",
-        strip_prefix = "rules_proto-11bf7c25e666dd7ddacbcd4d4c4a9de7a25175f8",
+        sha256 = "14a225870ab4e91869652cfd69ef2028277fc1dc4910d65d353b62d6e0ae21f4",
+        strip_prefix = "rules_proto-7.1.0",
         urls = tf_mirror_urls(
-            "https://github.com/bazelbuild/rules_proto/archive/11bf7c25e666dd7ddacbcd4d4c4a9de7a25175f8.tar.gz",
+            "https://github.com/bazelbuild/rules_proto/archive/refs/tags/7.1.0.tar.gz",
         ),
     )
 
diff --git a/third_party/xla/third_party/rules_proto.patch b/third_party/xla/third_party/rules_proto.patch
deleted file mode 100644
index 6868665a922075..00000000000000
--- a/third_party/xla/third_party/rules_proto.patch
+++ /dev/null
@@ -1,27 +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.
-# ==============================================================================
---- a/proto/private/native.bzl
-+++ b/proto/private/native.bzl
-@@ -26,7 +26,10 @@
- """Lovely workaround to be able to expose native constants pretending to be Starlark."""
-
-+load("@com_google_protobuf//bazel/common:proto_common.bzl", "proto_common")
-+load("@com_google_protobuf//bazel/common:proto_info.bzl", "ProtoInfo")
-+
- # buildifier: disable=native-proto
- NativeProtoInfo = ProtoInfo
-
- # buildifier: disable=native-proto
- native_proto_common = proto_common
\ No newline at end of file
diff --git a/third_party/xla/tools/def_file_filter/symbols_pybind.txt b/third_party/xla/tools/def_file_filter/symbols_pybind.txt
index 003a77aa2965b9..f8fce5e7322775 100644
--- a/third_party/xla/tools/def_file_filter/symbols_pybind.txt
+++ b/third_party/xla/tools/def_file_filter/symbols_pybind.txt
@@ -259,6 +259,7 @@ tensorflow::ImportFunction
 
 [//tensorflow/compiler/mlir/lite/python:converter_python_api] # converter_python_api
 tflite::Convert
+tflite::ConvertMlirBytecode
 tflite::MlirQuantizeModel
 tflite::MlirSparsifyModel
 tflite::RegisterCustomOpdefs
diff --git a/third_party/xla/workspace3.bzl b/third_party/xla/workspace3.bzl
index b326a9bc4a7938..79a81c7a644602 100644
--- a/third_party/xla/workspace3.bzl
+++ b/third_party/xla/workspace3.bzl
@@ -22,12 +22,11 @@ load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls")
 def workspace():
     tf_http_archive(
         name = "rules_proto",
-        sha256 = "20b240eba17a36be4b0b22635aca63053913d5c1ee36e16be36499d167a2f533",
-        strip_prefix = "rules_proto-11bf7c25e666dd7ddacbcd4d4c4a9de7a25175f8",
+        sha256 = "14a225870ab4e91869652cfd69ef2028277fc1dc4910d65d353b62d6e0ae21f4",
+        strip_prefix = "rules_proto-7.1.0",
         urls = tf_mirror_urls(
-            "https://github.com/bazelbuild/rules_proto/archive/11bf7c25e666dd7ddacbcd4d4c4a9de7a25175f8.tar.gz",
+            "https://github.com/bazelbuild/rules_proto/archive/refs/tags/7.1.0.tar.gz",
         ),
-        patch_file = ["//third_party:rules_proto.patch"],
     )
 
     # https://github.com/bazelbuild/bazel-skylib/releases
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 fb27b51af9de9f..4201819a1b030f 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/tests/BUILD
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/tests/BUILD
@@ -25,9 +25,18 @@ package(
     licenses = ["notice"],
 )
 
+_FILECHECK_ONLY_TESTS = [
+    "reduce_column_small/f32_32_v2.hlo",
+    "reduce_column_small/s8_f32_32_v4.hlo",
+    "scatter/sorted_indices_large.hlo",
+]
+
 lit_test_suite(
     name = "tests",
-    srcs = glob(["**/*.hlo"]),
+    srcs = glob(
+        ["**/*.hlo"],
+        exclude = _FILECHECK_ONLY_TESTS,
+    ),
     cfg = "//xla:lit.cfg.py",
     default_tags = tf_cuda_tests_tags(),
     exec_properties = tf_exec_properties({"tags": tf_cuda_tests_tags()}),
@@ -38,3 +47,15 @@ lit_test_suite(
         "@llvm-project//llvm:FileCheck",
     ],
 )
+
+lit_test_suite(
+    name = "filecheck_only_tests",
+    srcs = _FILECHECK_ONLY_TESTS,
+    cfg = "//xla:lit.cfg.py",
+    default_tags = ["gpu"],
+    tools = [
+        "//xla/backends/gpu/codegen/tools:fusion_to_mlir",
+        "//xla/codegen/tools:emitters_opt",
+        "@llvm-project//llvm:FileCheck",
+    ],
+)
diff --git a/third_party/xla/xla/backends/gpu/transforms/collectives/convert_async_collectives_to_sync_test.cc b/third_party/xla/xla/backends/gpu/transforms/collectives/convert_async_collectives_to_sync_test.cc
index a8138a37cadea8..a77354cb03eed2 100644
--- a/third_party/xla/xla/backends/gpu/transforms/collectives/convert_async_collectives_to_sync_test.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/collectives/convert_async_collectives_to_sync_test.cc
@@ -512,7 +512,7 @@ TEST_F(GpuConvertAsyncCollectivesToSyncTest,
     CHECK-NOT: all-reduce-start
     CHECK: %id2 = f32[] bitcast(%id)
     CHECK: ROOT %{{.*}} = u32[] all-reduce(%id)
-    CHECK-SAME: "is_sync":true,"is_pipelined":true
+    CHECK-SAME: "is_pipelined":true{{.*}}"is_sync":true
   )"),
               IsOkAndHolds(true));
 }
diff --git a/third_party/xla/xla/backends/gpu/transforms/convert_triton_gemm_config_test.cc b/third_party/xla/xla/backends/gpu/transforms/convert_triton_gemm_config_test.cc
index 8fc0976a7170ca..975f1430fb65e5 100644
--- a/third_party/xla/xla/backends/gpu/transforms/convert_triton_gemm_config_test.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/convert_triton_gemm_config_test.cc
@@ -100,12 +100,12 @@ ENTRY entry {
     CHECK: ENTRY
     CHECK: ROOT{{.*}}fusion(
     CHECK-SAME: kind=kCustom
-    CHECK-SAME: "kind":"__triton_nested_gemm_fusion"
     CHECK-SAME: "block_level_fusion_config"
-    CHECK-SAME: "num_warps":"4"
-    CHECK-SAME: "output_tiles":[{"sizes":["64","256"]}]
     CHECK-SAME: "num_ctas":3
     CHECK-SAME: "num_stages":5
+    CHECK-SAME: "num_warps":"4"
+    CHECK-SAME: "output_tiles":[{"sizes":["64","256"]}]
+    CHECK-SAME: "kind":"__triton_nested_gemm_fusion"
 )"));
   const HloInstruction* fusion = nullptr;
   ASSERT_THAT(module->entry_computation()->root_instruction(),
@@ -157,8 +157,8 @@ ENTRY entry {
     CHECK: ROOT {{.*}} = bf16[4,4]{1,0} scaled-dot({{.*}}backend_config={"sizes":["64"]}
     CHECK: ENTRY
     CHECK: ROOT{{.*}}fusion(
-    CHECK-SAME: "kind":"__triton_nested_gemm_fusion"
     CHECK-SAME: "output_tiles":[{"sizes":["16","32"]}]
+    CHECK-SAME: "kind":"__triton_nested_gemm_fusion"
 )"));
 }
 
diff --git a/third_party/xla/xla/backends/gpu/transforms/windowed_einsum_handler_test.cc b/third_party/xla/xla/backends/gpu/transforms/windowed_einsum_handler_test.cc
index e0a50c3c1ed5cd..f397673ea5fa85 100644
--- a/third_party/xla/xla/backends/gpu/transforms/windowed_einsum_handler_test.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/windowed_einsum_handler_test.cc
@@ -345,7 +345,7 @@ CHECK: }
 CHECK: dimensions={1}
 CHECK-DAG: %[[P0:.*]] = bf16[1,8192,32768]{2,1,0} parameter(0)
 CHECK-DAG: %[[SLICE4:.*]] = bf16[1,2048,32768]{2,1,0} slice(%[[P0:.*]]), slice={[0:1], [6144:8192], [0:32768]}
-CHECK-DAG: %[[DOT0:.*]] = bf16[1,4,2048,32768]{3,2,1,0} dot(%[[A2A0:.*]], %[[SLICE4:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"operation_queue_id":"8","force_earliest_schedule":false
+CHECK-DAG: %[[DOT0:.*]] = bf16[1,4,2048,32768]{3,2,1,0} dot(%[[A2A0:.*]], %[[SLICE4:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"device_type":"DEVICE_TYPE_INVALID","force_earliest_schedule":false,"operation_queue_id":"8"
 
 CHECK-DAG: %[[SLICE1:.*]] = bf16[1,4,2048,2048]{3,2,1,0} slice(%[[P1]]), slice={[0:1], [0:4], [0:2048], [4096:6144]}
 CHECK: %[[A2A1:.*]] = bf16[1,4,2048,2048]{3,2,1,0} all-to-all(%[[SLICE1]]),
@@ -354,7 +354,7 @@ CHECK:     {0,1,2,3},{4,5,6,7}
 CHECK: }
 CHECK: dimensions={1}
 CHECK-DAG: %[[SLICE5:.*]] = bf16[1,2048,32768]{2,1,0} slice(%[[P0:.*]]), slice={[0:1], [4096:6144], [0:32768]}
-CHECK-DAG: %[[DOT1:.*]] = bf16[1,4,2048,32768]{3,2,1,0} dot(%[[A2A1:.*]], %[[SLICE5:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"operation_queue_id":"7","force_earliest_schedule":false
+CHECK-DAG: %[[DOT1:.*]] = bf16[1,4,2048,32768]{3,2,1,0} dot(%[[A2A1:.*]], %[[SLICE5:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"device_type":"DEVICE_TYPE_INVALID","force_earliest_schedule":false,"operation_queue_id":"7"
 
 CHECK-DAG: %[[SLICE2:.*]] = bf16[1,4,2048,2048]{3,2,1,0} slice(%[[P1]]), slice={[0:1], [0:4], [0:2048], [2048:4096]}
 CHECK: %[[A2A2:.*]] = bf16[1,4,2048,2048]{3,2,1,0} all-to-all(%[[SLICE2]]),
@@ -363,7 +363,7 @@ CHECK:     {0,1,2,3},{4,5,6,7}
 CHECK: }
 CHECK: dimensions={1}
 CHECK-DAG: %[[SLICE6:.*]] = bf16[1,2048,32768]{2,1,0} slice(%[[P0:.*]]), slice={[0:1], [2048:4096], [0:32768]}
-CHECK-DAG: %[[DOT2:.*]] = bf16[1,4,2048,32768]{3,2,1,0} dot(%[[A2A2:.*]], %[[SLICE6:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"operation_queue_id":"6","force_earliest_schedule":false
+CHECK-DAG: %[[DOT2:.*]] = bf16[1,4,2048,32768]{3,2,1,0} dot(%[[A2A2:.*]], %[[SLICE6:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"device_type":"DEVICE_TYPE_INVALID","force_earliest_schedule":false,"operation_queue_id":"6"
 
 CHECK-DAG: %[[SLICE3:.*]] = bf16[1,4,2048,2048]{3,2,1,0} slice(%[[P1]]), slice={[0:1], [0:4], [0:2048], [0:2048]}
 CHECK: %[[A2A3:.*]] = bf16[1,4,2048,2048]{3,2,1,0} all-to-all(%[[SLICE3]]),
@@ -372,7 +372,7 @@ CHECK:     {0,1,2,3},{4,5,6,7}
 CHECK: }
 CHECK: dimensions={1}
 CHECK-DAG: %[[SLICE7:.*]] = bf16[1,2048,32768]{2,1,0} slice(%[[P0:.*]]), slice={[0:1], [0:2048], [0:32768]}
-CHECK-DAG: %[[DOT3:.*]] = bf16[1,4,2048,32768]{3,2,1,0} dot(%[[A2A3:.*]], %[[SLICE7:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"operation_queue_id":"5","force_earliest_schedule":false
+CHECK-DAG: %[[DOT3:.*]] = bf16[1,4,2048,32768]{3,2,1,0} dot(%[[A2A3:.*]], %[[SLICE7:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"device_type":"DEVICE_TYPE_INVALID","force_earliest_schedule":false,"operation_queue_id":"5"
 CHECK-DAG: %[[CONSTANT:.*]] = bf16[] constant(0)
 CHECK-DAG: %[[BROADCAST:.*]] = bf16[1,4,2048,32768]{3,2,1,0} broadcast(%[[CONSTANT:.*]]), dimensions={}
 CHECK-DAG: %[[ADD0:.*]] = bf16[1,4,2048,32768]{3,2,1,0} add(%[[DOT0:.*]], %[[BROADCAST:.*]])
@@ -415,7 +415,7 @@ CHECK-DAG: %[[P1:.*]] = bf16[1,4,2048,32768]{3,2,1,0} parameter(1)
 CHECK-DAG: %[[SLICE0:.*]] = bf16[1,4,2048,8192]{3,2,1,0} slice(%[[P1]]), slice={[0:1], [0:4], [0:2048], [24576:32768]}
 CHECK-DAG: %[[P0:.*]] = bf16[1,8192,32768]{2,1,0} parameter(0)
 CHECK-DAG: %[[SLICE4:.*]] = bf16[1,8192,8192]{2,1,0} slice(%[[P0:.*]]), slice={[0:1], [0:8192], [24576:32768]}
-CHECK-DAG: %[[DOT0:.*]] = bf16[1,4,2048,8192]{3,2,1,0} dot(%[[SLICE0:.*]], %[[SLICE4:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={2}, backend_config={"operation_queue_id":"8","force_earliest_schedule":false
+CHECK-DAG: %[[DOT0:.*]] = bf16[1,4,2048,8192]{3,2,1,0} dot(%[[SLICE0:.*]], %[[SLICE4:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={2}, backend_config={"device_type":"DEVICE_TYPE_INVALID","force_earliest_schedule":false,"operation_queue_id":"8"
 CHECK: %[[A2A0:.*]] = bf16[1,4,2048,8192]{3,2,1,0} all-to-all(%[[DOT0:.*]]),
 CHECK: replica_groups={
 CHECK:     {0,1,2,3}
@@ -424,7 +424,7 @@ CHECK: dimensions={1}
 
 CHECK-DAG: %[[SLICE1:.*]] = bf16[1,4,2048,8192]{3,2,1,0} slice(%[[P1]]), slice={[0:1], [0:4], [0:2048], [16384:24576]}
 CHECK-DAG: %[[SLICE5:.*]] = bf16[1,8192,8192]{2,1,0} slice(%[[P0:.*]]), slice={[0:1], [0:8192], [16384:24576]}
-CHECK-DAG: %[[DOT1:.*]] = bf16[1,4,2048,8192]{3,2,1,0} dot(%[[SLICE1:.*]], %[[SLICE5:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={2}, backend_config={"operation_queue_id":"7","force_earliest_schedule":false
+CHECK-DAG: %[[DOT1:.*]] = bf16[1,4,2048,8192]{3,2,1,0} dot(%[[SLICE1:.*]], %[[SLICE5:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={2}, backend_config={"device_type":"DEVICE_TYPE_INVALID","force_earliest_schedule":false,"operation_queue_id":"7"
 CHECK: %[[A2A1:.*]] = bf16[1,4,2048,8192]{3,2,1,0} all-to-all(%[[DOT1:.*]]),
 CHECK: replica_groups={
 CHECK:     {0,1,2,3}
@@ -433,7 +433,7 @@ CHECK: dimensions={1}
 
 CHECK-DAG: %[[SLICE2:.*]] = bf16[1,4,2048,8192]{3,2,1,0} slice(%[[P1]]), slice={[0:1], [0:4], [0:2048], [8192:16384]}
 CHECK-DAG: %[[SLICE6:.*]] = bf16[1,8192,8192]{2,1,0} slice(%[[P0:.*]]), slice={[0:1], [0:8192], [8192:16384]}
-CHECK-DAG: %[[DOT2:.*]] = bf16[1,4,2048,8192]{3,2,1,0} dot(%[[SLICE2:.*]], %[[SLICE6:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={2}, backend_config={"operation_queue_id":"6","force_earliest_schedule":false
+CHECK-DAG: %[[DOT2:.*]] = bf16[1,4,2048,8192]{3,2,1,0} dot(%[[SLICE2:.*]], %[[SLICE6:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={2}, backend_config={"device_type":"DEVICE_TYPE_INVALID","force_earliest_schedule":false,"operation_queue_id":"6"
 CHECK: %[[A2A2:.*]] = bf16[1,4,2048,8192]{3,2,1,0} all-to-all(%[[DOT2:.*]]),
 CHECK: replica_groups={
 CHECK:     {0,1,2,3}
@@ -442,7 +442,7 @@ CHECK: dimensions={1}
 
 CHECK-DAG: %[[SLICE3:.*]] = bf16[1,4,2048,8192]{3,2,1,0} slice(%[[P1]]), slice={[0:1], [0:4], [0:2048], [0:8192]}
 CHECK-DAG: %[[SLICE7:.*]] = bf16[1,8192,8192]{2,1,0} slice(%[[P0:.*]]), slice={[0:1], [0:8192], [0:8192]}
-CHECK-DAG: %[[DOT3:.*]] = bf16[1,4,2048,8192]{3,2,1,0} dot(%[[SLICE3:.*]], %[[SLICE7:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={2}, backend_config={"operation_queue_id":"5","force_earliest_schedule":false
+CHECK-DAG: %[[DOT3:.*]] = bf16[1,4,2048,8192]{3,2,1,0} dot(%[[SLICE3:.*]], %[[SLICE7:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={2}, backend_config={"device_type":"DEVICE_TYPE_INVALID","force_earliest_schedule":false,"operation_queue_id":"5"
 CHECK: %[[A2A3:.*]] = bf16[1,4,2048,8192]{3,2,1,0} all-to-all(%[[DOT3:.*]]),
 CHECK: replica_groups={
 CHECK:     {0,1,2,3}
@@ -505,7 +505,7 @@ CHECK: }
 CHECK: dimensions={1}
 CHECK-DAG: %[[P0:.*]] = bf16[1,8192,32768]{2,1,0} parameter(0)
 CHECK-DAG: %[[SLICE4:.*]] = bf16[1,2048,32768]{2,1,0} slice(%[[P0:.*]]), slice={[0:1], [6144:8192], [0:32768]}
-CHECK-DAG: %[[DOT0:.*]] = bf16[1,4,2048,32768]{3,2,1,0} dot(%[[A2A0:.*]], %[[SLICE4:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"operation_queue_id":"9","force_earliest_schedule":false
+CHECK-DAG: %[[DOT0:.*]] = bf16[1,4,2048,32768]{3,2,1,0} dot(%[[A2A0:.*]], %[[SLICE4:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"device_type":"DEVICE_TYPE_INVALID","force_earliest_schedule":false,"operation_queue_id":"9"
 
 CHECK-DAG: %[[SLICE1:.*]] = bf16[1,4,2048,2048]{3,2,1,0} slice(%[[COPY:.*]]), slice={[0:1], [0:4], [0:2048], [4096:6144]}
 CHECK: %[[A2A1:.*]] = bf16[1,4,2048,2048]{3,2,1,0} all-to-all(%[[SLICE1]]),
@@ -514,7 +514,7 @@ CHECK:     {0,1,2,3}
 CHECK: }
 CHECK: dimensions={1}
 CHECK-DAG: %[[SLICE5:.*]] = bf16[1,2048,32768]{2,1,0} slice(%[[P0:.*]]), slice={[0:1], [4096:6144], [0:32768]}
-CHECK-DAG: %[[DOT1:.*]] = bf16[1,4,2048,32768]{3,2,1,0} dot(%[[A2A1:.*]], %[[SLICE5:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"operation_queue_id":"8","force_earliest_schedule":false
+CHECK-DAG: %[[DOT1:.*]] = bf16[1,4,2048,32768]{3,2,1,0} dot(%[[A2A1:.*]], %[[SLICE5:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"device_type":"DEVICE_TYPE_INVALID","force_earliest_schedule":false,"operation_queue_id":"8"
 
 CHECK-DAG: %[[SLICE2:.*]] = bf16[1,4,2048,2048]{3,2,1,0} slice(%[[COPY:.*]]), slice={[0:1], [0:4], [0:2048], [2048:4096]}
 CHECK: %[[A2A2:.*]] = bf16[1,4,2048,2048]{3,2,1,0} all-to-all(%[[SLICE2]]),
@@ -523,7 +523,7 @@ CHECK:     {0,1,2,3}
 CHECK: }
 CHECK: dimensions={1}
 CHECK-DAG: %[[SLICE6:.*]] = bf16[1,2048,32768]{2,1,0} slice(%[[P0:.*]]), slice={[0:1], [2048:4096], [0:32768]}
-CHECK-DAG: %[[DOT2:.*]] = bf16[1,4,2048,32768]{3,2,1,0} dot(%[[A2A2:.*]], %[[SLICE6:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"operation_queue_id":"7","force_earliest_schedule":false
+CHECK-DAG: %[[DOT2:.*]] = bf16[1,4,2048,32768]{3,2,1,0} dot(%[[A2A2:.*]], %[[SLICE6:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"device_type":"DEVICE_TYPE_INVALID","force_earliest_schedule":false,"operation_queue_id":"7"
 
 CHECK-DAG: %[[SLICE3:.*]] = bf16[1,4,2048,2048]{3,2,1,0} slice(%[[COPY:.*]]), slice={[0:1], [0:4], [0:2048], [0:2048]}
 CHECK: %[[A2A3:.*]] = bf16[1,4,2048,2048]{3,2,1,0} all-to-all(%[[SLICE3]]),
@@ -532,7 +532,7 @@ CHECK:     {0,1,2,3}
 CHECK: }
 CHECK: dimensions={1}
 CHECK-DAG: %[[SLICE7:.*]] = bf16[1,2048,32768]{2,1,0} slice(%[[P0:.*]]), slice={[0:1], [0:2048], [0:32768]}
-CHECK-DAG: %[[DOT3:.*]] = bf16[1,4,2048,32768]{3,2,1,0} dot(%[[A2A3:.*]], %[[SLICE7:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"operation_queue_id":"6","force_earliest_schedule":false
+CHECK-DAG: %[[DOT3:.*]] = bf16[1,4,2048,32768]{3,2,1,0} dot(%[[A2A3:.*]], %[[SLICE7:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"device_type":"DEVICE_TYPE_INVALID","force_earliest_schedule":false,"operation_queue_id":"6"
 CHECK-DAG: %[[CONSTANT:.*]] = bf16[] constant(0)
 CHECK-DAG: %[[BROADCAST:.*]] = bf16[1,4,2048,32768]{3,2,1,0} broadcast(%[[CONSTANT:.*]]), dimensions={}
 CHECK-DAG: %[[ADD0:.*]] = bf16[1,4,2048,32768]{3,2,1,0} add(%[[DOT0:.*]], %[[BROADCAST:.*]])
@@ -581,7 +581,7 @@ CHECK-DAG: %[[P1:.*]] = bf16[1,4,2048,32768]{3,2,1,0} parameter(0)
 CHECK-DAG: %[[SLICE0:.*]] = bf16[1,4,2048,8192]{3,2,1,0} slice(%[[P1]]), slice={[0:1], [0:4], [0:2048], [24576:32768]}
 CHECK-DAG: %[[P0:.*]] = bf16[1,32768,8192]{2,1,0} parameter(1)
 CHECK-DAG: %[[SLICE4:.*]] = bf16[1,8192,8192]{2,1,0} slice(%[[P0:.*]]), slice={[0:1], [24576:32768], [0:8192]}
-CHECK-DAG: %[[DOT0:.*]] = bf16[1,4,2048,8192]{3,2,1,0} dot(%[[SLICE0:.*]], %[[SLICE4:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"operation_queue_id":"12","force_earliest_schedule":false
+CHECK-DAG: %[[DOT0:.*]] = bf16[1,4,2048,8192]{3,2,1,0} dot(%[[SLICE0:.*]], %[[SLICE4:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"device_type":"DEVICE_TYPE_INVALID","force_earliest_schedule":false,"operation_queue_id":"12"
 CHECK: %[[A2A0:.*]] = bf16[1,4,2048,8192]{3,2,1,0} all-to-all(%[[DOT0:.*]]),
 CHECK: replica_groups={
 CHECK:     {0,1,2,3}
@@ -590,7 +590,7 @@ CHECK: dimensions={1}
 
 CHECK-DAG: %[[SLICE1:.*]] = bf16[1,4,2048,8192]{3,2,1,0} slice(%[[P1]]), slice={[0:1], [0:4], [0:2048], [16384:24576]}
 CHECK-DAG: %[[SLICE5:.*]] = bf16[1,8192,8192]{2,1,0} slice(%[[P0:.*]]), slice={[0:1], [16384:24576], [0:8192]}
-CHECK-DAG: %[[DOT1:.*]] = bf16[1,4,2048,8192]{3,2,1,0} dot(%[[SLICE1:.*]], %[[SLICE5:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"operation_queue_id":"11","force_earliest_schedule":false
+CHECK-DAG: %[[DOT1:.*]] = bf16[1,4,2048,8192]{3,2,1,0} dot(%[[SLICE1:.*]], %[[SLICE5:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"device_type":"DEVICE_TYPE_INVALID","force_earliest_schedule":false,"operation_queue_id":"11"
 CHECK: %[[A2A1:.*]] = bf16[1,4,2048,8192]{3,2,1,0} all-to-all(%[[DOT1:.*]]),
 CHECK: replica_groups={
 CHECK:     {0,1,2,3}
@@ -599,7 +599,7 @@ CHECK: dimensions={1}
 
 CHECK-DAG: %[[SLICE2:.*]] = bf16[1,4,2048,8192]{3,2,1,0} slice(%[[P1]]), slice={[0:1], [0:4], [0:2048], [8192:16384]}
 CHECK-DAG: %[[SLICE6:.*]] = bf16[1,8192,8192]{2,1,0} slice(%[[P0:.*]]), slice={[0:1], [8192:16384], [0:8192]}
-CHECK-DAG: %[[DOT2:.*]] = bf16[1,4,2048,8192]{3,2,1,0} dot(%[[SLICE2:.*]], %[[SLICE6:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"operation_queue_id":"10","force_earliest_schedule":false
+CHECK-DAG: %[[DOT2:.*]] = bf16[1,4,2048,8192]{3,2,1,0} dot(%[[SLICE2:.*]], %[[SLICE6:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"device_type":"DEVICE_TYPE_INVALID","force_earliest_schedule":false,"operation_queue_id":"10"
 CHECK: %[[A2A2:.*]] = bf16[1,4,2048,8192]{3,2,1,0} all-to-all(%[[DOT2:.*]]),
 CHECK: replica_groups={
 CHECK:     {0,1,2,3}
@@ -608,7 +608,7 @@ CHECK: dimensions={1}
 
 CHECK-DAG: %[[SLICE3:.*]] = bf16[1,4,2048,8192]{3,2,1,0} slice(%[[P1]]), slice={[0:1], [0:4], [0:2048], [0:8192]}
 CHECK-DAG: %[[SLICE7:.*]] = bf16[1,8192,8192]{2,1,0} slice(%[[P0:.*]]), slice={[0:1], [0:8192], [0:8192]}
-CHECK-DAG: %[[DOT3:.*]] = bf16[1,4,2048,8192]{3,2,1,0} dot(%[[SLICE3:.*]], %[[SLICE7:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"operation_queue_id":"9","force_earliest_schedule":false
+CHECK-DAG: %[[DOT3:.*]] = bf16[1,4,2048,8192]{3,2,1,0} dot(%[[SLICE3:.*]], %[[SLICE7:.*]]), lhs_batch_dims={0}, lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}, backend_config={"device_type":"DEVICE_TYPE_INVALID","force_earliest_schedule":false,"operation_queue_id":"9"
 CHECK: %[[A2A3:.*]] = bf16[1,4,2048,8192]{3,2,1,0} all-to-all(%[[DOT3:.*]]),
 CHECK: replica_groups={
 CHECK:     {0,1,2,3}
diff --git a/third_party/xla/xla/backends/profiler/gpu/BUILD b/third_party/xla/xla/backends/profiler/gpu/BUILD
index 3bf88b2de50a18..7c7ea026cfe4a6 100644
--- a/third_party/xla/xla/backends/profiler/gpu/BUILD
+++ b/third_party/xla/xla/backends/profiler/gpu/BUILD
@@ -22,6 +22,10 @@ load(
     "if_google",
     "internal_visibility",
 )
+load(
+    "//xla/tsl:tsl.default.bzl",
+    "get_compatible_with_portable",
+)
 load("//xla/tsl/platform:rules_cc.bzl", "cc_library")
 load("//xla/tsl/platform/default:cuda_build_defs.bzl", "if_cuda_is_configured", "if_cuda_newer_than")
 
@@ -998,3 +1002,24 @@ cc_library(
         "@local_config_cuda//cuda:cuda_headers",
     ],
 )
+
+cc_library(
+    name = "cuda_graph_topology_mapper",
+    srcs = ["cuda_graph_topology_mapper.cc"],
+    hdrs = ["cuda_graph_topology_mapper.h"],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        "//xla/tsl/platform:logging",
+        "@com_google_absl//absl/container:flat_hash_map",
+    ],
+)
+
+xla_cc_test(
+    name = "cuda_graph_topology_mapper_test",
+    srcs = ["cuda_graph_topology_mapper_test.cc"],
+    deps = [
+        ":cuda_graph_topology_mapper",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
diff --git a/third_party/xla/xla/backends/profiler/gpu/cuda_graph_topology_mapper.cc b/third_party/xla/xla/backends/profiler/gpu/cuda_graph_topology_mapper.cc
new file mode 100644
index 00000000000000..5f5628e0b5d2c9
--- /dev/null
+++ b/third_party/xla/xla/backends/profiler/gpu/cuda_graph_topology_mapper.cc
@@ -0,0 +1,154 @@
+/* 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.
+==============================================================================*/
+
+#include "xla/backends/profiler/gpu/cuda_graph_topology_mapper.h"
+
+#include <cstddef>
+#include <cstdint>
+#include <utility>
+#include <vector>
+
+#include "absl/container/flat_hash_map.h"
+#include "xla/tsl/platform/logging.h"
+
+namespace xla {
+namespace profiler {
+
+size_t CudaGraphTopologyMapper::CalculateMergedSize(
+    uint32_t graph_id, const absl::flat_hash_map<uint32_t, size_t>& base_sizes,
+    const absl::flat_hash_map<uint32_t, std::vector<ChildGraphEntry>>&
+        child_graphs,
+    absl::flat_hash_map<uint32_t, size_t>* merged_sizes) {
+  DCHECK(merged_sizes != nullptr);
+  auto merged_it = merged_sizes->find(graph_id);
+  if (merged_it != merged_sizes->end()) {
+    return merged_it->second;
+  }
+
+  size_t size = 0;
+  auto size_it = base_sizes.find(graph_id);
+  if (size_it != base_sizes.end()) {
+    size = size_it->second;
+  }
+
+  auto children_it = child_graphs.find(graph_id);
+  if (children_it != child_graphs.end()) {
+    for (const auto& child : children_it->second) {
+      size_t child_size = CalculateMergedSize(child.child_graph_id, base_sizes,
+                                              child_graphs, merged_sizes);
+      size += child_size;
+      if (!child.is_conditional && size > 0) {
+        size -= 1;  // Child graph node itself is replaced for inline children.
+      }
+    }
+  }
+
+  (*merged_sizes)[graph_id] = size;
+  return size;
+}
+
+std::pair<uint32_t, uint32_t> CudaGraphTopologyMapper::ResolveMergedNode(
+    uint32_t graph_id, uint32_t node_index,
+    const absl::flat_hash_map<uint32_t, size_t>& base_sizes,
+    const absl::flat_hash_map<uint32_t, std::vector<ChildGraphEntry>>&
+        child_graphs,
+    absl::flat_hash_map<uint32_t, size_t>* merged_sizes) {
+  DCHECK(merged_sizes != nullptr);
+  auto children_it = child_graphs.find(graph_id);
+  if (children_it == child_graphs.end() || children_it->second.empty()) {
+    return {graph_id, node_index};
+  }
+
+  // Calculate total inline size.
+  size_t total_inline_size = 0;
+  auto size_it = base_sizes.find(graph_id);
+  if (size_it != base_sizes.end()) {
+    total_inline_size = size_it->second;
+  }
+  for (const auto& child : children_it->second) {
+    if (!child.is_conditional) {
+      size_t child_size = CalculateMergedSize(child.child_graph_id, base_sizes,
+                                              child_graphs, merged_sizes);
+      if (child_size > 0) {
+        total_inline_size += (child_size - 1);
+      } else if (total_inline_size > 0) {
+        total_inline_size -= 1;
+      }
+    }
+  }
+
+  if (node_index < total_inline_size) {
+    // 1. Resolve inline nodes
+    int64_t inline_offset = 0;
+    for (const auto& child : children_it->second) {
+      if (child.is_conditional) {
+        continue;
+      }
+      size_t child_size = CalculateMergedSize(child.child_graph_id, base_sizes,
+                                              child_graphs, merged_sizes);
+      int64_t child_start =
+          static_cast<int64_t>(child.insertion_point) + inline_offset;
+
+      if (child_size > 0) {
+        int64_t child_end = child_start + static_cast<int64_t>(child_size) - 1;
+        if (static_cast<int64_t>(node_index) >= child_start &&
+            static_cast<int64_t>(node_index) <= child_end) {
+          uint32_t local_index = static_cast<uint32_t>(
+              static_cast<int64_t>(node_index) - child_start);
+          return ResolveMergedNode(child.child_graph_id, local_index,
+                                   base_sizes, child_graphs, merged_sizes);
+        }
+      }
+
+      if (static_cast<int64_t>(node_index) < child_start) {
+        return {graph_id,
+                static_cast<uint32_t>(static_cast<int64_t>(node_index) -
+                                      inline_offset)};
+      }
+
+      inline_offset += (static_cast<int64_t>(child_size) - 1);
+    }
+    return {graph_id, static_cast<uint32_t>(static_cast<int64_t>(node_index) -
+                                            inline_offset)};
+  }
+
+  // 2. Resolve conditional nodes
+  uint32_t conditional_offset = 0;
+  for (const auto& child : children_it->second) {
+    if (!child.is_conditional) {
+      continue;
+    }
+    size_t child_size = CalculateMergedSize(child.child_graph_id, base_sizes,
+                                            child_graphs, merged_sizes);
+    if (child_size == 0) {
+      continue;
+    }
+    uint32_t child_start = total_inline_size + conditional_offset;
+    uint32_t child_end = child_start + child_size - 1;
+
+    if (node_index >= child_start && node_index <= child_end) {
+      uint32_t local_index = node_index - child_start;
+      return ResolveMergedNode(child.child_graph_id, local_index, base_sizes,
+                               child_graphs, merged_sizes);
+    }
+
+    conditional_offset += child_size;
+  }
+
+  return {graph_id, node_index};
+}
+
+}  // namespace profiler
+}  // namespace xla
diff --git a/third_party/xla/xla/backends/profiler/gpu/cuda_graph_topology_mapper.h b/third_party/xla/xla/backends/profiler/gpu/cuda_graph_topology_mapper.h
new file mode 100644
index 00000000000000..6fab464bd06b89
--- /dev/null
+++ b/third_party/xla/xla/backends/profiler/gpu/cuda_graph_topology_mapper.h
@@ -0,0 +1,98 @@
+/* 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.
+==============================================================================*/
+
+#ifndef XLA_BACKENDS_PROFILER_GPU_CUDA_GRAPH_TOPOLOGY_MAPPER_H_
+#define XLA_BACKENDS_PROFILER_GPU_CUDA_GRAPH_TOPOLOGY_MAPPER_H_
+
+#include <cstddef>
+#include <cstdint>
+#include <utility>
+#include <vector>
+
+#include "absl/container/flat_hash_map.h"
+
+namespace xla {
+namespace profiler {
+
+// CudaGraphTopologyMapper is a pure, state-less C++ helper library designed to
+// mirror the layout and node-flattening sequences that the NVIDIA CUDA driver
+// uses to assign executable flat LocalNodeId sequences to CUPTI activity traces
+// during the instantiation of hierarchical and conditional nested CUDA Graphs.
+//
+// Topology Examples:
+//
+// 1. Standard Inline Child Graph (`is_conditional = false`)
+// In this scenario, the nodes of the nested child graph completely replace the
+// placeholder node in the parent graph.
+// Layout Example:
+//   Parent Graph (Size 3): [P0, P_Child, P2]
+//   Child Graph (Size 2): [C0, C1]
+//   Expected Flattened Execution sequence: P0 -> C0 -> C1 -> P2
+//   Expected Flat indices: P0 (0), C0 (1), C1 (2), P2 (3)
+//
+// 2. Conditional Nested Graph (`is_conditional = true`)
+// In this scenario, the child graph's nodes are executed as a separate entity
+// at runtime, appended to the end of the parent graph's sequence. The
+// placeholder node in the parent retains its original index in the flattened
+// space.
+// Layout Example:
+//   Parent Graph (Size 3): [P0, P1_Cond, P2]
+//   Child Graph (Size 2): [C0, C1]
+//   Expected Flattened Sequence (Kernels only): P0 -> C0 -> C1 -> P2
+//   Expected Flat indices: P0 (0), P1_Cond (1), P2 (2), C0 (3), C1 (4)
+//   (Note that the CUPTI trace for `C0` reports local node index `3`, while the
+//    original placeholder `P2` reports index `2`).
+//
+// All APIs in this library are pure, static, and have zero CUDA SDK
+// dependencies, allowing them to run identically on the host CPU for unit
+// testing.
+class CudaGraphTopologyMapper {
+ public:
+  struct ChildGraphEntry {
+    uint32_t child_graph_id;
+    uint32_t insertion_point;
+    bool is_conditional;
+
+    bool operator<(const ChildGraphEntry& other) const {
+      return insertion_point < other.insertion_point;
+    }
+  };
+
+  // Recursively calculates the total merged size of a graph, caching results in
+  // the provided mutable merged_sizes map.
+  static size_t CalculateMergedSize(
+      uint32_t graph_id,
+      const absl::flat_hash_map<uint32_t, size_t>& base_sizes,
+      const absl::flat_hash_map<uint32_t, std::vector<ChildGraphEntry>>&
+          child_graphs,
+      absl::flat_hash_map<uint32_t, size_t>* merged_sizes);
+
+  // Resolves a flat, merged node index into its target (template_graph_id,
+  // local_node_index) pair.
+  static std::pair<uint32_t, uint32_t> ResolveMergedNode(
+      uint32_t graph_id, uint32_t node_index,
+      const absl::flat_hash_map<uint32_t, size_t>& base_sizes,
+      const absl::flat_hash_map<uint32_t, std::vector<ChildGraphEntry>>&
+          child_graphs,
+      absl::flat_hash_map<uint32_t, size_t>* merged_sizes);
+
+ private:
+  CudaGraphTopologyMapper() = delete;
+};
+
+}  // namespace profiler
+}  // namespace xla
+
+#endif  // XLA_BACKENDS_PROFILER_GPU_CUDA_GRAPH_TOPOLOGY_MAPPER_H_
diff --git a/third_party/xla/xla/backends/profiler/gpu/cuda_graph_topology_mapper_test.cc b/third_party/xla/xla/backends/profiler/gpu/cuda_graph_topology_mapper_test.cc
new file mode 100644
index 00000000000000..1a6b506bdcc6a0
--- /dev/null
+++ b/third_party/xla/xla/backends/profiler/gpu/cuda_graph_topology_mapper_test.cc
@@ -0,0 +1,338 @@
+/* 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.
+==============================================================================*/
+
+#include "xla/backends/profiler/gpu/cuda_graph_topology_mapper.h"
+
+#include <cstddef>
+#include <cstdint>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "absl/container/flat_hash_map.h"
+
+namespace xla {
+namespace profiler {
+namespace {
+
+TEST(CudaGraphTopologyMapperTest, CalculateMergedSizeSimpleInline) {
+  absl::flat_hash_map<uint32_t, size_t> base_sizes = {{1, 3}, {2, 2}};
+  absl::flat_hash_map<uint32_t,
+                      std::vector<CudaGraphTopologyMapper::ChildGraphEntry>>
+      child_graphs;
+  child_graphs[1].push_back({2, 1, /*is_conditional=*/false});
+
+  absl::flat_hash_map<uint32_t, size_t> merged_sizes;
+  size_t size = CudaGraphTopologyMapper::CalculateMergedSize(
+      1, base_sizes, child_graphs, &merged_sizes);
+
+  // 3 + 2 - 1 = 4
+  EXPECT_EQ(size, 4);
+}
+
+TEST(CudaGraphTopologyMapperTest, ResolveMergedNodeInline) {
+  absl::flat_hash_map<uint32_t, size_t> base_sizes = {{1, 3}, {2, 2}};
+  absl::flat_hash_map<uint32_t,
+                      std::vector<CudaGraphTopologyMapper::ChildGraphEntry>>
+      child_graphs;
+  child_graphs[1].push_back({2, 1, /*is_conditional=*/false});
+
+  absl::flat_hash_map<uint32_t, size_t> merged_sizes;
+
+  // Expected Merged indices:
+  // 0 -> Parent node 0
+  // 1 -> Child node 0
+  // 2 -> Child node 1
+  // 3 -> Parent node 2
+  auto res0 = CudaGraphTopologyMapper::ResolveMergedNode(
+      1, 0, base_sizes, child_graphs, &merged_sizes);
+  EXPECT_EQ(res0.first, 1);
+  EXPECT_EQ(res0.second, 0);
+
+  auto res1 = CudaGraphTopologyMapper::ResolveMergedNode(
+      1, 1, base_sizes, child_graphs, &merged_sizes);
+  EXPECT_EQ(res1.first, 2);
+  EXPECT_EQ(res1.second, 0);
+
+  auto res2 = CudaGraphTopologyMapper::ResolveMergedNode(
+      1, 2, base_sizes, child_graphs, &merged_sizes);
+  EXPECT_EQ(res2.first, 2);
+  EXPECT_EQ(res2.second, 1);
+
+  auto res3 = CudaGraphTopologyMapper::ResolveMergedNode(
+      1, 3, base_sizes, child_graphs, &merged_sizes);
+  EXPECT_EQ(res3.first, 1);
+  EXPECT_EQ(res3.second, 2);
+}
+
+TEST(CudaGraphTopologyMapperTest, ResolveMergedNodeDeepNestingInline) {
+  // Parent (1) -> Child (2) -> Grandchild (3)
+  absl::flat_hash_map<uint32_t, size_t> base_sizes = {{1, 3}, {2, 3}, {3, 2}};
+  absl::flat_hash_map<uint32_t,
+                      std::vector<CudaGraphTopologyMapper::ChildGraphEntry>>
+      child_graphs;
+  child_graphs[1].push_back({2, 1, /*is_conditional=*/false});
+  child_graphs[2].push_back({3, 1, /*is_conditional=*/false});
+
+  absl::flat_hash_map<uint32_t, size_t> merged_sizes;
+
+  // Expected indices:
+  // 0 -> Parent 0
+  // 1 -> Child 0
+  // 2 -> Grandchild 0
+  // 3 -> Grandchild 1
+  // 4 -> Child 2
+  // 5 -> Parent 2
+  auto res2 = CudaGraphTopologyMapper::ResolveMergedNode(
+      1, 2, base_sizes, child_graphs, &merged_sizes);
+  EXPECT_EQ(res2.first, 3);
+  EXPECT_EQ(res2.second, 0);
+
+  auto res5 = CudaGraphTopologyMapper::ResolveMergedNode(
+      1, 5, base_sizes, child_graphs, &merged_sizes);
+  EXPECT_EQ(res5.first, 1);
+  EXPECT_EQ(res5.second, 2);
+}
+
+TEST(CudaGraphTopologyMapperTest, CalculateMergedSizeConditional) {
+  absl::flat_hash_map<uint32_t, size_t> base_sizes = {{1, 3}, {2, 2}};
+  absl::flat_hash_map<uint32_t,
+                      std::vector<CudaGraphTopologyMapper::ChildGraphEntry>>
+      child_graphs;
+  child_graphs[1].push_back({2, 1, /*is_conditional=*/true});
+
+  absl::flat_hash_map<uint32_t, size_t> merged_sizes;
+  size_t size = CudaGraphTopologyMapper::CalculateMergedSize(
+      1, base_sizes, child_graphs, &merged_sizes);
+
+  // 3 + 2 = 5 (Conditional child graph nodes are appended, size is not shifted
+  // by -1)
+  EXPECT_EQ(size, 5);
+}
+
+TEST(CudaGraphTopologyMapperTest, ResolveMergedNodeConditional) {
+  absl::flat_hash_map<uint32_t, size_t> base_sizes = {{1, 3}, {2, 2}};
+  absl::flat_hash_map<uint32_t,
+                      std::vector<CudaGraphTopologyMapper::ChildGraphEntry>>
+      child_graphs;
+  child_graphs[1].push_back({2, 1, /*is_conditional=*/true});
+
+  absl::flat_hash_map<uint32_t, size_t> merged_sizes;
+
+  // Expected Merged indices for Conditional Nested Graph (Hardware Ground
+  // Truth): 0 -> Parent node 0 1 -> Parent node 1 2 -> Parent node 2 3 -> Child
+  // node 0 4 -> Child node 1
+  auto res0 = CudaGraphTopologyMapper::ResolveMergedNode(
+      1, 0, base_sizes, child_graphs, &merged_sizes);
+  EXPECT_EQ(res0.first, 1);
+  EXPECT_EQ(res0.second, 0);
+
+  auto res1 = CudaGraphTopologyMapper::ResolveMergedNode(
+      1, 1, base_sizes, child_graphs, &merged_sizes);
+  EXPECT_EQ(res1.first, 1);
+  EXPECT_EQ(res1.second, 1);
+
+  auto res2 = CudaGraphTopologyMapper::ResolveMergedNode(
+      1, 2, base_sizes, child_graphs, &merged_sizes);
+  EXPECT_EQ(res2.first, 1);
+  EXPECT_EQ(res2.second, 2);
+
+  auto res3 = CudaGraphTopologyMapper::ResolveMergedNode(
+      1, 3, base_sizes, child_graphs, &merged_sizes);
+  EXPECT_EQ(res3.first, 2);
+  EXPECT_EQ(res3.second, 0);
+
+  auto res4 = CudaGraphTopologyMapper::ResolveMergedNode(
+      1, 4, base_sizes, child_graphs, &merged_sizes);
+  EXPECT_EQ(res4.first, 2);
+  EXPECT_EQ(res4.second, 1);
+}
+
+TEST(CudaGraphTopologyMapperTest, ResolveMergedNodeMixedInlineAndConditional) {
+  absl::flat_hash_map<uint32_t, size_t> base_sizes = {{1, 3}, {2, 2}, {3, 2}};
+  absl::flat_hash_map<uint32_t,
+                      std::vector<CudaGraphTopologyMapper::ChildGraphEntry>>
+      child_graphs;
+  child_graphs[1].push_back({2, 1, /*is_conditional=*/false});
+  child_graphs[1].push_back({3, 2, /*is_conditional=*/true});
+
+  absl::flat_hash_map<uint32_t, size_t> merged_sizes;
+
+  auto res = CudaGraphTopologyMapper::ResolveMergedNode(
+      1, 4, base_sizes, child_graphs, &merged_sizes);
+  EXPECT_EQ(res.first, 3);
+  EXPECT_EQ(res.second, 0);
+}
+
+TEST(CudaGraphTopologyMapperTest, ResolveMergedNodeMultipleConditional) {
+  absl::flat_hash_map<uint32_t, size_t> base_sizes = {{1, 3}, {2, 2}, {3, 2}};
+  absl::flat_hash_map<uint32_t,
+                      std::vector<CudaGraphTopologyMapper::ChildGraphEntry>>
+      child_graphs;
+  child_graphs[1].push_back({2, 1, /*is_conditional=*/true});
+  child_graphs[1].push_back({3, 2, /*is_conditional=*/true});
+
+  absl::flat_hash_map<uint32_t, size_t> merged_sizes;
+
+  auto res = CudaGraphTopologyMapper::ResolveMergedNode(
+      1, 5, base_sizes, child_graphs, &merged_sizes);
+  EXPECT_EQ(res.first, 3);
+  EXPECT_EQ(res.second, 0);
+}
+
+TEST(CudaGraphTopologyMapperTest, ResolveMergedNodeOutOfBounds) {
+  absl::flat_hash_map<uint32_t, size_t> base_sizes = {{1, 3}, {2, 2}};
+  absl::flat_hash_map<uint32_t,
+                      std::vector<CudaGraphTopologyMapper::ChildGraphEntry>>
+      child_graphs;
+  child_graphs[1].push_back({2, 1, /*is_conditional=*/false});
+
+  absl::flat_hash_map<uint32_t, size_t> merged_sizes;
+
+  auto res = CudaGraphTopologyMapper::ResolveMergedNode(
+      1, 4, base_sizes, child_graphs, &merged_sizes);
+  EXPECT_EQ(res.first, 1);
+  EXPECT_EQ(res.second, 4);
+}
+
+TEST(CudaGraphTopologyMapperTest, EmptyChildGraphHandling) {
+  // Test 1: Inline empty child graph (child size == 0).
+  // Parent (1) size 2, Child (2) size 0.
+  // Merged size should be: 2 + 0 - 1 = 1.
+  {
+    absl::flat_hash_map<uint32_t, size_t> base_sizes = {{1, 2}, {2, 0}};
+    absl::flat_hash_map<uint32_t,
+                        std::vector<CudaGraphTopologyMapper::ChildGraphEntry>>
+        child_graphs;
+    child_graphs[1].push_back({2, 1, /*is_conditional=*/false});
+
+    absl::flat_hash_map<uint32_t, size_t> merged_sizes;
+    size_t size = CudaGraphTopologyMapper::CalculateMergedSize(
+        1, base_sizes, child_graphs, &merged_sizes);
+    EXPECT_EQ(size, 1);
+
+    auto res0 = CudaGraphTopologyMapper::ResolveMergedNode(
+        1, 0, base_sizes, child_graphs, &merged_sizes);
+    EXPECT_EQ(res0.first, 1);
+    EXPECT_EQ(res0.second, 0);
+  }
+
+  // Test 2: Conditional empty child graph (child size == 0).
+  // Parent (1) size 2, Child (2) size 0.
+  // Merged size should be: 2 + 0 = 2.
+  {
+    absl::flat_hash_map<uint32_t, size_t> base_sizes = {{1, 2}, {2, 0}};
+    absl::flat_hash_map<uint32_t,
+                        std::vector<CudaGraphTopologyMapper::ChildGraphEntry>>
+        child_graphs;
+    child_graphs[1].push_back({2, 1, /*is_conditional=*/true});
+
+    absl::flat_hash_map<uint32_t, size_t> merged_sizes;
+    size_t size = CudaGraphTopologyMapper::CalculateMergedSize(
+        1, base_sizes, child_graphs, &merged_sizes);
+    EXPECT_EQ(size, 2);
+
+    auto res0 = CudaGraphTopologyMapper::ResolveMergedNode(
+        1, 0, base_sizes, child_graphs, &merged_sizes);
+    EXPECT_EQ(res0.first, 1);
+    EXPECT_EQ(res0.second, 0);
+
+    auto res1 = CudaGraphTopologyMapper::ResolveMergedNode(
+        1, 1, base_sizes, child_graphs, &merged_sizes);
+    EXPECT_EQ(res1.first, 1);
+    EXPECT_EQ(res1.second, 1);
+  }
+
+  // Test 3: Completely empty parent and child graphs.
+  {
+    absl::flat_hash_map<uint32_t, size_t> base_sizes = {{1, 0}, {2, 0}};
+    absl::flat_hash_map<uint32_t,
+                        std::vector<CudaGraphTopologyMapper::ChildGraphEntry>>
+        child_graphs;
+    child_graphs[1].push_back({2, 0, /*is_conditional=*/false});
+
+    absl::flat_hash_map<uint32_t, size_t> merged_sizes;
+    size_t size = CudaGraphTopologyMapper::CalculateMergedSize(
+        1, base_sizes, child_graphs, &merged_sizes);
+    EXPECT_EQ(size, 0);
+  }
+}
+
+TEST(CudaGraphTopologyMapperTest, EmptyChildGraphResolutionWithTrailingNodes) {
+  // 1. Parent size 3 (nodes P0, P1, P2) with empty inline child at insertion
+  // point 1:
+  {
+    absl::flat_hash_map<uint32_t, size_t> base_sizes = {{1, 3}, {2, 0}};
+    absl::flat_hash_map<uint32_t,
+                        std::vector<CudaGraphTopologyMapper::ChildGraphEntry>>
+        child_graphs;
+    child_graphs[1].push_back({2, 1, /*is_conditional=*/false});
+
+    absl::flat_hash_map<uint32_t, size_t> merged_sizes;
+    size_t size = CudaGraphTopologyMapper::CalculateMergedSize(
+        1, base_sizes, child_graphs, &merged_sizes);
+    EXPECT_EQ(size, 2);
+
+    auto res0 = CudaGraphTopologyMapper::ResolveMergedNode(
+        1, 0, base_sizes, child_graphs, &merged_sizes);
+    EXPECT_EQ(res0.first, 1);
+    EXPECT_EQ(res0.second, 0);
+
+    auto res1 = CudaGraphTopologyMapper::ResolveMergedNode(
+        1, 1, base_sizes, child_graphs, &merged_sizes);
+    EXPECT_EQ(res1.first, 1);
+    EXPECT_EQ(res1.second, 2);
+  }
+
+  // 2. Parent size 3 with empty inline child at insertion point 1 AND
+  // conditional child (ID 3, size 2):
+  {
+    absl::flat_hash_map<uint32_t, size_t> base_sizes = {
+        {{1, 3}, {2, 0}, {3, 2}}};
+    absl::flat_hash_map<uint32_t,
+                        std::vector<CudaGraphTopologyMapper::ChildGraphEntry>>
+        child_graphs;
+    child_graphs[1].push_back({2, 1, /*is_conditional=*/false});
+    child_graphs[1].push_back({3, 2, /*is_conditional=*/true});
+
+    absl::flat_hash_map<uint32_t, size_t> merged_sizes;
+    size_t size = CudaGraphTopologyMapper::CalculateMergedSize(
+        1, base_sizes, child_graphs, &merged_sizes);
+    EXPECT_EQ(size, 4);
+
+    auto res0 = CudaGraphTopologyMapper::ResolveMergedNode(
+        1, 0, base_sizes, child_graphs, &merged_sizes);
+    EXPECT_EQ(res0.first, 1);
+    EXPECT_EQ(res0.second, 0);
+
+    auto res1 = CudaGraphTopologyMapper::ResolveMergedNode(
+        1, 1, base_sizes, child_graphs, &merged_sizes);
+    EXPECT_EQ(res1.first, 1);
+    EXPECT_EQ(res1.second, 2);
+
+    auto res2 = CudaGraphTopologyMapper::ResolveMergedNode(
+        1, 2, base_sizes, child_graphs, &merged_sizes);
+    EXPECT_EQ(res2.first, 3);
+    EXPECT_EQ(res2.second, 0);
+
+    auto res3 = CudaGraphTopologyMapper::ResolveMergedNode(
+        1, 3, base_sizes, child_graphs, &merged_sizes);
+    EXPECT_EQ(res3.first, 3);
+    EXPECT_EQ(res3.second, 1);
+  }
+}
+
+}  // namespace
+}  // namespace profiler
+}  // namespace xla
diff --git a/third_party/xla/xla/benchmarks/README.md b/third_party/xla/xla/benchmarks/README.md
new file mode 100644
index 00000000000000..64d2b447423de9
--- /dev/null
+++ b/third_party/xla/xla/benchmarks/README.md
@@ -0,0 +1,21 @@
+# TPU Microbenchmarks
+
+## Quickstart
+
+From the root directory of the XLA project, run:
+```bash
+# Run setup script to create the venv and install dependencies
+./xla/benchmarks/setup.sh
+
+# Activate the venv
+source xla/benchmarks/.venv/bin/activate
+
+# Run an individual benchmark
+python3 xla/benchmarks/pallas_microbenchmarks/dense_matmul.py  --dim=1,2048,2048,2048 --fmt=f8e4m3fn,f8e4m3fn,f32
+
+# Run dense matmul benchmark suite and write results to a CSV file
+python3 xla/benchmarks/run_benchmarks.py --benchmarks=dense_matmul --csv_path=<path_to_dir>
+
+# Run full benchmark suite and write results to multiple CSV files
+python3 xla/benchmarks/run_benchmarks.py --csv_path=<path_to_dir>
+```
\ No newline at end of file
diff --git a/third_party/xla/xla/benchmarks/core/benchmark.py b/third_party/xla/xla/benchmarks/core/benchmark.py
index e85f992707ab60..7889f2c05fd5eb 100644
--- a/third_party/xla/xla/benchmarks/core/benchmark.py
+++ b/third_party/xla/xla/benchmarks/core/benchmark.py
@@ -238,14 +238,22 @@ def run(
     return profiler_results
 
 
-@dataclasses.dataclass(frozen=True)
+@dataclasses.dataclass(frozen=True, repr=False)
 class BenchmarkConfig(abc.ABC):
   """Base class for benchmark configs."""
 
   def as_dict(self) -> dict[str, Any]:
     """Returns a dictionary representation of the config."""
+
+    def _is_dtype(val):
+      try:
+        jnp.dtype(val)
+        return True
+      except (TypeError, ValueError):
+        return False
+
     return {
-        k: dtype_to_str(v) if isinstance(v, jnp.dtype) else v
+        k: dtype_to_str(v) if _is_dtype(v) else v
         for k, v in dataclasses.asdict(self).items()
     }
 
diff --git a/third_party/xla/xla/benchmarks/jax_microbenchmarks/matmul_lib.py b/third_party/xla/xla/benchmarks/jax_microbenchmarks/matmul_lib.py
index f10dd1e4987572..7b736c874a20cf 100644
--- a/third_party/xla/xla/benchmarks/jax_microbenchmarks/matmul_lib.py
+++ b/third_party/xla/xla/benchmarks/jax_microbenchmarks/matmul_lib.py
@@ -24,7 +24,7 @@
 from xla.benchmarks.core import benchmark
 
 
-@dataclasses.dataclass(frozen=True, kw_only=True)
+@dataclasses.dataclass(frozen=True, kw_only=True, repr=False)
 class JaxMatmulConfig(benchmark.BenchmarkConfig):
   """Config for JAX matmul benchmark.
 
diff --git a/third_party/xla/xla/benchmarks/pallas_microbenchmarks/dense_matmul_lib.py b/third_party/xla/xla/benchmarks/pallas_microbenchmarks/dense_matmul_lib.py
index fa988e6e4bdd9b..7b4cd08420df62 100644
--- a/third_party/xla/xla/benchmarks/pallas_microbenchmarks/dense_matmul_lib.py
+++ b/third_party/xla/xla/benchmarks/pallas_microbenchmarks/dense_matmul_lib.py
@@ -94,7 +94,7 @@ def select_window(
   return int(block_m), int(block_k), int(block_n)
 
 
-@dataclasses.dataclass(frozen=True, kw_only=True)
+@dataclasses.dataclass(frozen=True, kw_only=True, repr=False)
 class DenseMatmulConfig(benchmark.BenchmarkConfig):
   """Config for Pallas dense matmul benchmark.
 
diff --git a/third_party/xla/xla/benchmarks/pallas_microbenchmarks/subchannel_matmul_lib.py b/third_party/xla/xla/benchmarks/pallas_microbenchmarks/subchannel_matmul_lib.py
index 20838543385c2d..8a8f29fbbe6cdb 100644
--- a/third_party/xla/xla/benchmarks/pallas_microbenchmarks/subchannel_matmul_lib.py
+++ b/third_party/xla/xla/benchmarks/pallas_microbenchmarks/subchannel_matmul_lib.py
@@ -125,7 +125,7 @@ def select_window(
   return int(block_m), int(block_k), int(block_n)
 
 
-@dataclasses.dataclass(frozen=True, kw_only=True)
+@dataclasses.dataclass(frozen=True, kw_only=True, repr=False)
 class SubchannelMatmulConfig(benchmark.BenchmarkConfig):
   """Config for Pallas subchannel quantized matmul benchmark.
 
diff --git a/third_party/xla/xla/benchmarks/requirements.txt b/third_party/xla/xla/benchmarks/requirements.txt
new file mode 100644
index 00000000000000..77971c82529003
--- /dev/null
+++ b/third_party/xla/xla/benchmarks/requirements.txt
@@ -0,0 +1,20 @@
+# 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.
+
+-f https://storage.googleapis.com/jax-releases/libtpu_releases.html
+absl-py
+immutabledict
+jax[tpu]
+numpy
+pandas
diff --git a/third_party/xla/xla/benchmarks/setup.sh b/third_party/xla/xla/benchmarks/setup.sh
new file mode 100755
index 00000000000000..63fe0b17e8f708
--- /dev/null
+++ b/third_party/xla/xla/benchmarks/setup.sh
@@ -0,0 +1,77 @@
+#!/bin/bash
+# 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.
+# ==============================================================================
+
+set -e
+set -u
+set -o pipefail
+
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &> /dev/null && pwd)"
+REPO_ROOT="$(cd -- "${SCRIPT_DIR}/../.." &> /dev/null && pwd)"
+VENV_DIR="${1:-${VENV_DIR:-${SCRIPT_DIR}/.venv}}"
+
+echo "============================================================"
+echo "Setting up environment for OpenXLA TPU microbenchmarks"
+echo "Repository root: ${REPO_ROOT}"
+echo "Virtual environment directory: ${VENV_DIR}"
+echo "============================================================"
+
+# Check for python3
+if ! command -v python3 &> /dev/null; then
+  echo "Error: python3 is not installed or not found in PATH." >&2
+  exit 1
+fi
+
+# Create virtual environment if it does not already exist
+if [[ ! -d "${VENV_DIR}" ]]; then
+  echo "Creating virtual environment at ${VENV_DIR}..."
+  python3 -m venv "${VENV_DIR}"
+else
+  echo "Using existing virtual environment at ${VENV_DIR}."
+fi
+
+# Activate virtual environment
+# shellcheck source=/dev/null
+source "${VENV_DIR}/bin/activate"
+
+# Upgrade pip
+echo "Upgrading pip..."
+pip install --upgrade pip
+
+# Install dependencies
+echo "Installing requirements from ${SCRIPT_DIR}/requirements.txt..."
+pip install -r "${SCRIPT_DIR}/requirements.txt"
+
+# Configure site-packages .pth file so xla is importable from anywhere
+echo "Configuring Python site-packages path for OpenXLA repository..."
+python3 -c "
+import site, pathlib, sys
+site_dirs = site.getsitepackages()
+if site_dirs:
+    pth_file = pathlib.Path(site_dirs[0]) / 'openxla_benchmarks.pth'
+    pth_file.write_text(sys.argv[1] + '\n')
+" "${REPO_ROOT}"
+
+export PYTHONPATH="${REPO_ROOT}:${PYTHONPATH:-}"
+
+echo "============================================================"
+echo "Setup completed successfully!"
+echo ""
+echo "To activate this virtual environment in your shell, run:"
+echo "  source ${VENV_DIR}/bin/activate"
+echo ""
+echo "To run the benchmarks suite, run:"
+echo "  python3 ${SCRIPT_DIR}/run_benchmarks.py"
+echo "============================================================"
diff --git a/third_party/xla/xla/codegen/emitters/tests/BUILD b/third_party/xla/xla/codegen/emitters/tests/BUILD
index a1f04cceb1584c..7cedec6f546e7c 100644
--- a/third_party/xla/xla/codegen/emitters/tests/BUILD
+++ b/third_party/xla/xla/codegen/emitters/tests/BUILD
@@ -44,9 +44,35 @@ copy_file(
     is_executable = True,
 )
 
+#Need no gpu to run
+_FILECHECK_ONLY_TESTS = [
+    #TODO(rocm) Most of thease test cuda lowering only
+    "loop/acos_f32.hlo",
+    "loop/acos_f64.hlo",
+    "loop/acosh_f32.hlo",
+    "loop/acosh_f64.hlo",
+    "loop/asin_f32.hlo",
+    "loop/asin_f64.hlo",
+    "loop/asinh_f32.hlo",
+    "loop/asinh_f64.hlo",
+    "loop/atanh_f32.hlo",
+    "loop/atanh_f64.hlo",
+    "loop/broadcast_constant_block_dim_limit.hlo",
+    "loop/cosh_f32.hlo",
+    "loop/cosh_f64.hlo",
+    "loop/dot_fp8.hlo",
+    "loop/large_loop_slow_compile_time.hlo",
+    "loop/sinh_f32.hlo",
+    "loop/sinh_f64.hlo",
+    "concatenate/test_small_dim.hlo",
+]
+
 lit_test_suite(
     name = "tests",
-    srcs = glob(["**/*.hlo"]),
+    srcs = glob(
+        ["**/*.hlo"],
+        exclude = _FILECHECK_ONLY_TESTS,
+    ),
     cfg = "//xla:lit.cfg.py",
     default_tags = tf_cuda_tests_tags(),
     exec_properties = tf_exec_properties({"tags": tf_cuda_tests_tags()}),
@@ -59,3 +85,16 @@ lit_test_suite(
         "@llvm-project//llvm:FileCheck",
     ],
 )
+
+lit_test_suite(
+    name = "filecheck_only_tests",
+    srcs = _FILECHECK_ONLY_TESTS,
+    cfg = "//xla:lit.cfg.py",
+    default_tags = ["gpu"],
+    tools = [
+        ":cpu_fusion_to_mlir_copy",
+        ":gpu_fusion_to_mlir_copy",
+        "//xla/codegen/tools:emitters_opt",
+        "@llvm-project//llvm:FileCheck",
+    ],
+)
diff --git a/third_party/xla/xla/debug_options_flags.cc b/third_party/xla/xla/debug_options_flags.cc
index 6f8b72260dd115..66a1e7eb3767c5 100644
--- a/third_party/xla/xla/debug_options_flags.cc
+++ b/third_party/xla/xla/debug_options_flags.cc
@@ -265,7 +265,7 @@ DebugOptions DefaultDebugOptionsIgnoringFlags() {
   opts.set_xla_cpu_use_acl(true);
 #endif
   opts.set_xla_cpu_use_xnnpack(true);
-  opts.set_xla_cpu_use_new_xtile_lowering(false);
+  opts.set_xla_cpu_use_new_xtile_lowering(true);
   opts.set_xla_cpu_experimental_xnn_graph_fusion_mode(
       DebugOptions::XNN_GRAPH_FUSION_MODE_DISABLED);
   opts.add_xla_cpu_experimental_ynn_fusion_type(
diff --git a/third_party/xla/xla/hlo/ir/hlo_module.cc b/third_party/xla/xla/hlo/ir/hlo_module.cc
index fb20e46c044cfa..c8fa7c566d2db2 100644
--- a/third_party/xla/xla/hlo/ir/hlo_module.cc
+++ b/third_party/xla/xla/hlo/ir/hlo_module.cc
@@ -580,6 +580,7 @@ std::string HloModule::ToString() const {
   print_options.set_print_inline_stack_frames(
       db_options.xla_hlo_print_inline_stack_frames());
   print_options.set_compact_gte(db_options.xla_dump_compact_gte());
+  print_options.set_sort_backend_config(true);
   return ToString(print_options);
 }
 
diff --git a/third_party/xla/xla/hlo/ir/hlo_module_test.cc b/third_party/xla/xla/hlo/ir/hlo_module_test.cc
index b5cc565d3d8f93..e747795bfd523a 100644
--- a/third_party/xla/xla/hlo/ir/hlo_module_test.cc
+++ b/third_party/xla/xla/hlo/ir/hlo_module_test.cc
@@ -776,6 +776,22 @@ TEST(HloModuleTest, CheckToStringHonorsDebugOptions) {
   EXPECT_TRUE(filecheck_matched);
 }
 
+TEST(HloModuleTest, CheckToStringSortsBackendConfig) {
+  const char* hlo = R"(
+  HloModule test
+
+  ENTRY main {
+    ROOT custom-call = () custom-call(), custom_call_target="test", backend_config={"tuning_knobs":{"3":"0","2":"2"}}
+  })";
+
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnUnverifiedModule(hlo));
+  EXPECT_THAT(
+      module->ToString(),
+      ::testing::HasSubstr(
+          R"json(backend_config={"tuning_knobs":{"2":"2","3":"0"}})json"));
+}
+
 TEST(HloModuleTest, TestCallersAndCallees) {
   const char* hlo = R"(
     HloModule jit_h
diff --git a/third_party/xla/xla/service/spmd/gather_scatter_handler.cc b/third_party/xla/xla/service/spmd/gather_scatter_handler.cc
index 64c6834351ccd1..9d22a72e8d29a3 100644
--- a/third_party/xla/xla/service/spmd/gather_scatter_handler.cc
+++ b/third_party/xla/xla/service/spmd/gather_scatter_handler.cc
@@ -977,6 +977,7 @@ absl::StatusOr<HloInstruction*> PartitionGather(
           output_shape, operand.Replicate().hlo(), indices.Replicate().hlo(),
           gather->gather_dimension_numbers(), slice_sizes,
           gather->indices_are_sorted()));
+  new_gather->set_frontend_attributes(gather->frontend_attributes());
   new_gather->set_sharding(HloSharding::Replicate());
   new_gather = PartitionedHlo(new_gather, new_gather->shape(), operand.state())
                    .Reshard(output_sharding)
@@ -1020,6 +1021,7 @@ absl::Status SpmdPartitioningVisitor::HandleGatherWithoutConflicts(
         builder()->AddInstruction(HloInstruction::CreateGather(
             pshape, operand.hlo(), indices.hlo(), dnums, pslice_sizes,
             gather->indices_are_sorted()));
+    phlo->set_frontend_attributes(gather->frontend_attributes());
 
     SetPartitionedHlo(hlo, phlo);
     return absl::OkStatus();
@@ -1047,6 +1049,7 @@ absl::Status SpmdPartitioningVisitor::HandleGatherWithoutConflicts(
   HloInstruction* pgather = b->AddInstruction(HloInstruction::CreateGather(
       pshape, operand.hlo(), adjusted_indices_hlo, dnums, pslice_sizes,
       gather->indices_are_sorted()));
+  pgather->set_frontend_attributes(gather->frontend_attributes());
 
   const Shape filter_shape =
       ShapeUtil::ChangeElementType(indices.hlo()->shape(), PRED);
@@ -2007,6 +2010,7 @@ absl::StatusOr<HloInstruction*> PartitionScatter(
               scatter->to_apply()->Clone()),
           scatter->scatter_dimension_numbers(), scatter->indices_are_sorted(),
           scatter->unique_indices()));
+  new_scatter->set_frontend_attributes(scatter->frontend_attributes());
   new_scatter->set_sharding(
       HloSharding::Replicate().NormalizeTupleSharding(new_scatter->shape()));
   new_scatter =
@@ -2123,6 +2127,7 @@ absl::Status SpmdPartitioningVisitor::HandleScatterWithoutConflicts(
           scatter->to_apply()->Clone()),
       scatter->scatter_dimension_numbers(), scatter->indices_are_sorted(),
       scatter->unique_indices()));
+  pscatter->set_frontend_attributes(scatter->frontend_attributes());
   pscatter->set_sharding(HloSharding::Single(
       pscatter->shape(), hlo->sharding().IsTuple()
                              ? hlo->sharding().tuple_elements()[0]
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 38b62c9357248f..4f264492bc84d8 100644
--- a/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc
+++ b/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc
@@ -10626,6 +10626,132 @@ ENTRY entry {
   }
 }
 
+// `_xla_compute_type` selects SparseCore offload and is read long after SPMD,
+// so a partitioned gather/scatter must carry it across.
+// `need_resolve_conflicts` picks between two entirely separate implementations,
+// and each builds its own replacement instruction, so both are exercised.
+namespace {
+
+void ExpectComputeTypeAttr(const HloInstruction* instr) {
+  ASSERT_NE(instr, nullptr);
+  std::optional<std::string> attr =
+      instr->get_frontend_attribute("_xla_compute_type");
+  ASSERT_TRUE(attr.has_value());
+  EXPECT_EQ(*attr, "sparseoffload");
+}
+
+}  // namespace
+
+TEST_P(SpmdPartitioningTest, PassthroughGatherPreservesFrontendAttributes) {
+  absl::string_view hlo_string = R"(
+HloModule module
+
+ENTRY entry {
+  %input = f32[2,9] parameter(0), sharding={devices=[1,2]<=[2]}
+  %indices = s32[3] parameter(1), sharding={replicated}
+  ROOT %gather = f32[3,9] gather(%input, %indices), offset_dims={1},
+    collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1,
+    slice_sizes={1,9}, sharding={devices=[1,2]<=[2]},
+    frontend_attributes={_xla_compute_type="sparseoffload"}
+})";
+  for (bool need_resolve_conflicts : {true, false}) {
+    SpmdPartitionerOptions options;
+    options.need_resolve_conflicts = need_resolve_conflicts;
+    ASSERT_OK_AND_ASSIGN(
+        auto module,
+        PartitionComputation(hlo_string, /*num_devices=*/2, options));
+    ExpectComputeTypeAttr(FindInstruction(module.get(), HloOpcode::kGather));
+  }
+}
+
+TEST_P(SpmdPartitioningTest,
+       GatherPartitionedOnTrivialSliceDimsPreservesFrontendAttributes) {
+  absl::string_view hlo_string = R"(
+HloModule module
+
+ENTRY entry {
+  %input = f32[17,9] parameter(0), sharding={devices=[2,1]<=[2]}
+  %indices = s32[2,3] parameter(1), sharding={replicated}
+  ROOT %gather = f32[2,3,9] gather(%input, %indices), offset_dims={2},
+    collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=2,
+    slice_sizes={1,9}, sharding={replicated},
+    frontend_attributes={_xla_compute_type="sparseoffload"}
+})";
+  for (bool need_resolve_conflicts : {true, false}) {
+    SpmdPartitionerOptions options;
+    options.need_resolve_conflicts = need_resolve_conflicts;
+    ASSERT_OK_AND_ASSIGN(
+        auto module,
+        PartitionComputation(hlo_string, /*num_devices=*/2, options));
+    ExpectComputeTypeAttr(FindInstruction(module.get(), HloOpcode::kGather));
+  }
+}
+
+TEST_P(SpmdPartitioningTest, PassthroughScatterPreservesFrontendAttributes) {
+  absl::string_view hlo_string = R"(
+HloModule module
+
+add (lhs: f32[], rhs: f32[]) -> f32[] {
+  lhs = f32[] parameter(0)
+  rhs = f32[] parameter(1)
+  ROOT sum = f32[] add(lhs, rhs)
+}
+
+ENTRY entry {
+  %input = f32[2,9] parameter(0), sharding={devices=[1,2]<=[2]}
+  %indices = s32[3] parameter(1), sharding={replicated}
+  %updates = f32[3,9] parameter(2), sharding={devices=[1,2]<=[2]}
+  ROOT %scatter = f32[2,9] scatter(%input, %indices, %updates),
+      to_apply=add,
+      update_window_dims={1},
+      inserted_window_dims={0},
+      scatter_dims_to_operand_dims={0},
+      index_vector_dim=1, sharding={devices=[1,2]<=[2]},
+      frontend_attributes={_xla_compute_type="sparseoffload"}
+})";
+  for (bool need_resolve_conflicts : {true, false}) {
+    SpmdPartitionerOptions options;
+    options.need_resolve_conflicts = need_resolve_conflicts;
+    ASSERT_OK_AND_ASSIGN(
+        auto module,
+        PartitionComputation(hlo_string, /*num_devices=*/2, options));
+    ExpectComputeTypeAttr(FindInstruction(module.get(), HloOpcode::kScatter));
+  }
+}
+
+TEST_P(SpmdPartitioningTest,
+       ScatterPartitionedOnTrivialSliceDimsPreservesFrontendAttributes) {
+  absl::string_view hlo_string = R"(
+HloModule module
+
+add (lhs: f32[], rhs: f32[]) -> f32[] {
+  lhs = f32[] parameter(0)
+  rhs = f32[] parameter(1)
+  ROOT sum = f32[] add(lhs, rhs)
+}
+
+ENTRY entry {
+  %input = f32[17,9] parameter(0), sharding={devices=[2,1]<=[2]}
+  %indices = s32[2,3] parameter(1), sharding={replicated}
+  %updates = f32[2,3,9] parameter(2), sharding={replicated}
+  ROOT %scatter = f32[17,9] scatter(%input, %indices, %updates),
+      to_apply=add,
+      update_window_dims={2},
+      inserted_window_dims={0},
+      scatter_dims_to_operand_dims={0},
+      index_vector_dim=2, sharding={devices=[2,1]<=[2]},
+      frontend_attributes={_xla_compute_type="sparseoffload"}
+})";
+  for (bool need_resolve_conflicts : {true, false}) {
+    SpmdPartitionerOptions options;
+    options.need_resolve_conflicts = need_resolve_conflicts;
+    ASSERT_OK_AND_ASSIGN(
+        auto module,
+        PartitionComputation(hlo_string, /*num_devices=*/2, options));
+    ExpectComputeTypeAttr(FindInstruction(module.get(), HloOpcode::kScatter));
+  }
+}
+
 TEST_P(SpmdPartitioningTest,
        ScatterPartitionedOnTrivialSliceDims_PartialReplicate) {
   absl::string_view hlo_string = R"(
diff --git a/third_party/xla/xla/sort_json_test.cc b/third_party/xla/xla/sort_json_test.cc
index c24d5ab689b40d..cdbd3567c23325 100644
--- a/third_party/xla/xla/sort_json_test.cc
+++ b/third_party/xla/xla/sort_json_test.cc
@@ -46,5 +46,11 @@ TEST(SortJsonTest, SortsJson) {
               absl_testing::IsOkAndHolds(R"({"a":"a","a":"a}"})"));
 }
 
+TEST(SortJsonTest, SortsTuningKnobs) {
+  EXPECT_THAT(SortJson(R"({"algorithm":{"tuning_knobs":{"3":"0","2":"2"}}})"),
+              absl_testing::IsOkAndHolds(
+                  R"({"algorithm":{"tuning_knobs":{"2":"2","3":"0"}}})"));
+}
+
 }  // namespace
 }  // namespace xla
diff --git a/third_party/xla/xla/tools/cost_model/BUILD b/third_party/xla/xla/tools/cost_model/BUILD
index 5b3e54eef7c5ad..fe1e2715bbdbff 100644
--- a/third_party/xla/xla/tools/cost_model/BUILD
+++ b/third_party/xla/xla/tools/cost_model/BUILD
@@ -3,6 +3,10 @@ load(
     "//xla:xla.default.bzl",
     "xla_cc_test",
 )
+load(
+    "//xla/tsl/platform/default:cuda_build_defs.bzl",
+    "if_cuda_is_configured",
+)
 
 package(
     # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
@@ -15,17 +19,32 @@ cc_library(
     srcs = ["gpu_bandwidth_benchmark.cc"],
     hdrs = ["gpu_bandwidth_benchmark.h"],
     deps = [
+        "//xla/stream_executor:device_description",
+        "//xla/stream_executor:platform",
+        "//xla/stream_executor:platform_manager",
+        "//xla/stream_executor/gpu:gpu_init",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_macros",
+        "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings:str_format",
         "@com_google_absl//absl/types:span",
-    ],
+    ] + if_cuda_is_configured([
+        "//xla/stream_executor:cuda_platform",
+    ]),
 )
 
 xla_cc_test(
     name = "gpu_bandwidth_benchmark_test",
     srcs = ["gpu_bandwidth_benchmark_test.cc"],
+    tags = [
+        "cuda-only",
+        "requires-gpu-nvidia",
+    ],
     deps = [
         ":gpu_bandwidth_benchmark",
         "//xla/tests:xla_internal_test_main",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_matchers",
         "@com_google_googletest//:gtest",
     ],
 )
diff --git a/third_party/xla/xla/tools/cost_model/gpu_bandwidth_benchmark.cc b/third_party/xla/xla/tools/cost_model/gpu_bandwidth_benchmark.cc
index 9ce9311bd2a4fe..2ef01e08c97b7c 100644
--- a/third_party/xla/xla/tools/cost_model/gpu_bandwidth_benchmark.cc
+++ b/third_party/xla/xla/tools/cost_model/gpu_bandwidth_benchmark.cc
@@ -15,13 +15,39 @@ limitations under the License.
 
 #include "xla/tools/cost_model/gpu_bandwidth_benchmark.h"
 
+#include <cstdint>
+#include <memory>
 #include <string>
 
+#include "absl/status/status.h"
+#include "absl/status/status_macros.h"
+#include "absl/status/statusor.h"
 #include "absl/strings/str_format.h"
 #include "absl/types/span.h"
+#include "xla/stream_executor/device_description.h"
+#include "xla/stream_executor/gpu/gpu_init.h"
+#include "xla/stream_executor/platform.h"
+#include "xla/stream_executor/platform_manager.h"
 
 namespace xla::gpu {
 
+absl::StatusOr<double> GetPeakBandwidthBytesPerSec(int device_id) {
+  ABSL_ASSIGN_OR_RETURN(stream_executor::Platform * platform,
+                   stream_executor::PlatformManager::PlatformWithName(
+                       stream_executor::GpuPlatformName()));
+  ABSL_ASSIGN_OR_RETURN(
+      std::unique_ptr<stream_executor::DeviceDescription> description,
+      platform->DescriptionForDevice(device_id));
+  const int64_t bandwidth = description->memory_bandwidth();
+  if (bandwidth <= 0) {
+    return absl::InternalError(absl::StrFormat(
+        "Failed to determine peak memory bandwidth for device %d: "
+        "memory_bandwidth is %v.",
+        device_id, bandwidth));
+  }
+  return static_cast<double>(bandwidth);
+}
+
 std::string FormatBandwidthTable(absl::Span<const BandwidthEntry> entries) {
   std::string result =
       "DMA Size (Bytes)    Bandwidth Fraction\n"
diff --git a/third_party/xla/xla/tools/cost_model/gpu_bandwidth_benchmark.h b/third_party/xla/xla/tools/cost_model/gpu_bandwidth_benchmark.h
index b6842b9e50623d..aa4daa0acc05b5 100644
--- a/third_party/xla/xla/tools/cost_model/gpu_bandwidth_benchmark.h
+++ b/third_party/xla/xla/tools/cost_model/gpu_bandwidth_benchmark.h
@@ -19,6 +19,7 @@ limitations under the License.
 #include <cstdint>
 #include <string>
 
+#include "absl/status/statusor.h"
 #include "absl/types/span.h"
 
 namespace xla::gpu {
@@ -30,6 +31,10 @@ struct BandwidthEntry {
   float bandwidth_fraction = 0.0f;
 };
 
+// Returns theoretical peak GPU memory bandwidth in bytes per second for
+// `device_id`.
+absl::StatusOr<double> GetPeakBandwidthBytesPerSec(int device_id);
+
 // Formats bandwidth table entries into a human-readable table.
 std::string FormatBandwidthTable(absl::Span<const BandwidthEntry> entries);
 
diff --git a/third_party/xla/xla/tools/cost_model/gpu_bandwidth_benchmark_test.cc b/third_party/xla/xla/tools/cost_model/gpu_bandwidth_benchmark_test.cc
index 4d70629374f67b..eafd6f7d3dc574 100644
--- a/third_party/xla/xla/tools/cost_model/gpu_bandwidth_benchmark_test.cc
+++ b/third_party/xla/xla/tools/cost_model/gpu_bandwidth_benchmark_test.cc
@@ -17,11 +17,19 @@ limitations under the License.
 
 #include <vector>
 
+#include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include "absl/status/status.h"
+#include "absl/status/status_matchers.h"
 
 namespace xla::gpu {
 namespace {
 
+using ::absl_testing::IsOk;
+using ::absl_testing::IsOkAndHolds;
+using ::testing::Gt;
+using ::testing::Not;
+
 TEST(GpuBandwidthBenchmarkTest, FormatBandwidthTableEmpty) {
   EXPECT_EQ(FormatBandwidthTable({}),
             "DMA Size (Bytes)    Bandwidth Fraction\n"
@@ -29,7 +37,7 @@ TEST(GpuBandwidthBenchmarkTest, FormatBandwidthTableEmpty) {
 }
 
 TEST(GpuBandwidthBenchmarkTest, FormatBandwidthTableMultipleEntries) {
-  std::vector<BandwidthEntry> entries = {
+  const std::vector<BandwidthEntry> entries = {
       {8192, 0.00043418f},
       {16384, 0.00092645f},
       {32768, 0.00184066f},
@@ -44,5 +52,19 @@ TEST(GpuBandwidthBenchmarkTest, FormatBandwidthTableMultipleEntries) {
             "      8589934592            1.00000000\n");
 }
 
+TEST(GpuBandwidthBenchmarkTest, GetPeakBandwidthValidDevice) {
+  const absl::StatusOr<double> peak_bw =
+      GetPeakBandwidthBytesPerSec(/*device_id=*/0);
+  if (!peak_bw.ok()) {
+    GTEST_SKIP() << "No GPU device available: " << peak_bw.status();
+  }
+  EXPECT_THAT(peak_bw, IsOkAndHolds(Gt(0.0)));
+}
+
+TEST(GpuBandwidthBenchmarkTest, GetPeakBandwidthInvalidDevice) {
+  EXPECT_THAT(GetPeakBandwidthBytesPerSec(/*device_id=*/-1), Not(IsOk()));
+  EXPECT_THAT(GetPeakBandwidthBytesPerSec(/*device_id=*/9999), Not(IsOk()));
+}
+
 }  // namespace
 }  // namespace xla::gpu