diff --git a/tensorflow/compiler/jit/device_compilation_profiler.cc b/tensorflow/compiler/jit/device_compilation_profiler.cc
index 3af7856a226459..5976c7996c7263 100644
--- a/tensorflow/compiler/jit/device_compilation_profiler.cc
+++ b/tensorflow/compiler/jit/device_compilation_profiler.cc
@@ -96,8 +96,8 @@ void DeviceCompilationProfiler::RegisterExecution(
 
 absl::Status DeviceCompilationProfiler::RegisterCompilation(
     const NameAttrList& function, int64_t compile_time_us,
-    bool used_persistent_cache) {
-  metrics::UpdateXlaCompilationTime(compile_time_us);
+    bool used_persistent_cache, int64_t compile_end_us) {
+  metrics::UpdateXlaCompilationTime(compile_time_us, compile_end_us);
 
   const std::string& function_name = function.name();
 
diff --git a/tensorflow/compiler/jit/device_compilation_profiler.h b/tensorflow/compiler/jit/device_compilation_profiler.h
index 9f1d9521f4f1be..c58382c9286377 100644
--- a/tensorflow/compiler/jit/device_compilation_profiler.h
+++ b/tensorflow/compiler/jit/device_compilation_profiler.h
@@ -76,7 +76,8 @@ class DeviceCompilationProfiler : public ResourceBase {
   // XlaJitCompilationActivity.
   virtual absl::Status RegisterCompilation(const NameAttrList& function,
                                            int64_t compile_time_us,
-                                           bool used_persistent_cache);
+                                           bool used_persistent_cache,
+                                           int64_t compile_end_us = 0);
 
   void IncrementOngoingAsyncCompilations();
   void DecrementOngoingAsyncCompilations();
diff --git a/tensorflow/compiler/jit/device_compiler.h b/tensorflow/compiler/jit/device_compiler.h
index 3381895ee3eee2..2b424a9bbce729 100644
--- a/tensorflow/compiler/jit/device_compiler.h
+++ b/tensorflow/compiler/jit/device_compiler.h
@@ -332,6 +332,7 @@ DeviceCompiler<ExecutableType, ClientType>::CompileStrict(
     DeviceCompilationProfiler* profiler, mutex* mu) {
   tensorflow::Env* env = tensorflow::Env::Default();
   const uint64_t compile_start_us = env->NowMicros();
+  metrics::UpdateXlaCompilationStartTime(compile_start_us);
 
   TfGraphToHloCompiler compiler(options);
   cache_value.compile_state = DeviceCompileState::kCompiled;
@@ -390,7 +391,8 @@ DeviceCompiler<ExecutableType, ClientType>::CompileStrict(
 
   device_compiler_internal::LogOnceXlaCompiledFirstCluster();
   TF_RETURN_IF_ERROR(profiler->RegisterCompilation(
-      function, compile_time_us, loaded_executable.has_value()));
+      function, compile_time_us, loaded_executable.has_value(),
+      compile_end_us));
   return cache_value;
 }
 
diff --git a/tensorflow/compiler/jit/device_compiler_test.cc b/tensorflow/compiler/jit/device_compiler_test.cc
index 9e2adf614bc5a4..d42f797526be83 100644
--- a/tensorflow/compiler/jit/device_compiler_test.cc
+++ b/tensorflow/compiler/jit/device_compiler_test.cc
@@ -156,7 +156,7 @@ class MockDeviceCompilationProfiler : public DeviceCompilationProfiler {
               (override));
   MOCK_METHOD(absl::Status, RegisterCompilation,
               (const NameAttrList& function, int64_t compile_time_us,
-               bool used_persistent_cache),
+               bool used_persistent_cache, int64_t compile_end_us),
               (override));
 };
 
@@ -313,7 +313,7 @@ TEST_F(DeviceCompilerTest, CompileAsyncSuccess) {
   EXPECT_CALL(*mock_profiler_,
               ShouldCompileCluster(_, DeviceCompileMode::kAsync, 1))
       .WillOnce(Return(true));
-  EXPECT_CALL(*mock_profiler_, RegisterCompilation(_, _, false))
+  EXPECT_CALL(*mock_profiler_, RegisterCompilation(_, _, false, _))
       .WillOnce([&done] {
         done.Notify();
         return absl::OkStatus();
diff --git a/tensorflow/compiler/mlir/lite/BUILD b/tensorflow/compiler/mlir/lite/BUILD
index f9e50f01ab7481..26fbb92a7e3fb1 100644
--- a/tensorflow/compiler/mlir/lite/BUILD
+++ b/tensorflow/compiler/mlir/lite/BUILD
@@ -370,6 +370,35 @@ cc_library(
     ],
 )
 
+tf_cc_binary(
+    name = "litert-reduce",
+    testonly = True,
+    deps = [
+        ":litert_mlir_reduce_main",
+    ],
+)
+
+cc_library(
+    name = "litert_mlir_reduce_main",
+    testonly = True,
+    srcs = ["litert_mlir_reduce_main.cc"],
+    deps = [
+        ":lift_tflite_flex_ops",  # buildcleaner:keep
+        ":register_lite_dialects",
+        ":tensorflow_lite",
+        ":tf_tfl_passes",  # buildcleaner:keep
+        "//tensorflow/compiler/mlir:init_mlir",
+        "//tensorflow/compiler/mlir:passes",
+        "//tensorflow/compiler/mlir:register_common_dialects",
+        "//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_passes",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:AllPassesAndDialects",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:MlirReduceLib",
+    ],
+    alwayslink = 1,
+)
+
 cc_library(
     name = "utils",
     hdrs = ["utils/utils.h"],
@@ -1786,6 +1815,7 @@ cc_library(
         "//tensorflow/core:framework",
         "//tensorflow/core:portable_gif_internal",
         "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/platform:status",
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:flat_hash_map",
@@ -1797,6 +1827,7 @@ cc_library(
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:str_format",
+        "@eigen_archive//:eigen3",
         "@flatbuffers",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:ArithDialect",
diff --git a/tensorflow/compiler/mlir/lite/converter_flags.proto b/tensorflow/compiler/mlir/lite/converter_flags.proto
index cba7892e26f5e9..dfd2e956f1ca59 100644
--- a/tensorflow/compiler/mlir/lite/converter_flags.proto
+++ b/tensorflow/compiler/mlir/lite/converter_flags.proto
@@ -39,7 +39,7 @@ enum FileFormat {
 // of as properties of models, instead describing how models are to be
 // processed in the context of the present tooling job.
 //
-// Next ID to use: 72.
+// Next ID to use: 74.
 message ConverterFlags {
   reserved 54, 61;
 
@@ -390,4 +390,12 @@ message ConverterFlags {
   // If true, fold 16-bit float (fp16/bf16) to 32-bit float (fp32) casts on
   // large resource constants.
   optional bool fold_fp16_resource_casts = 71 [default = true];
+
+  // If true, enable debugging facilities such as IR cloning and file
+  // serialization.
+  optional bool enable_debug = 72 [default = false];
+
+  // Directory for debug artifacts and printed IR dumps when debug mode is
+  // enabled.
+  optional string debug_dir = 73;
 }
diff --git a/tensorflow/compiler/mlir/lite/debug/BUILD b/tensorflow/compiler/mlir/lite/debug/BUILD
index 14fe586abe2a72..50f6f6b6824c71 100644
--- a/tensorflow/compiler/mlir/lite/debug/BUILD
+++ b/tensorflow/compiler/mlir/lite/debug/BUILD
@@ -25,7 +25,7 @@ cc_library(
     deps = [
         ":debug_options_proto_cc",
         "//tensorflow/compiler/mlir/lite/metrics:error_collector_inst",
-        "//tensorflow/core:portable_gif_internal",
+        "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/status",
@@ -40,7 +40,6 @@ cc_library(
         "@llvm-project//mlir:Support",
         "@llvm-project//mlir:Transforms",
         "@tsl//tsl/platform:path",
-        "@tsl//tsl/platform:stringpiece",
         "@xla//xla/tsl/lib/io:buffered_file",
         "@xla//xla/tsl/platform:env",
     ],
diff --git a/tensorflow/compiler/mlir/lite/debug/debug.cc b/tensorflow/compiler/mlir/lite/debug/debug.cc
index 3b5ec4e4a84cad..f52aec086415b8 100644
--- a/tensorflow/compiler/mlir/lite/debug/debug.cc
+++ b/tensorflow/compiler/mlir/lite/debug/debug.cc
@@ -17,13 +17,18 @@ limitations under the License.
 
 #include <stddef.h>
 #include <stdint.h>
+#if defined(__linux__)
+#include <unistd.h>
+#endif
 
+#include <fstream>
 #include <functional>
 #include <memory>
 #include <string>
 #include <utility>
 #include <vector>
 
+#include "absl/container/flat_hash_map.h"
 #include "absl/container/flat_hash_set.h"
 #include "absl/log/log.h"
 #include "absl/status/status.h"
@@ -52,9 +57,7 @@ limitations under the License.
 #include "xla/tsl/lib/io/buffered_file.h"
 #include "xla/tsl/platform/env.h"
 #include "xla/tsl/platform/file_system.h"
-#include "tensorflow/core/platform/logging.h"
 #include "tsl/platform/path.h"
-#include "tsl/platform/stringpiece.h"
 
 // IWYU pragma: no_include "util/regexp/re2/re2.h"
 
@@ -166,12 +169,14 @@ std::string Sanitize(absl::string_view string) {
 // instrumentation dumps MLIR to external directories for convenience.
 class DumpInstrumentation : public mlir::PassInstrumentation {
  public:
-  explicit DumpInstrumentation(absl::string_view dump_dir,
-                               absl::string_view dump_pass_regex,
-                               absl::string_view dump_func_regex)
+  explicit DumpInstrumentation(
+      absl::string_view dump_dir, absl::string_view dump_pass_regex,
+      absl::string_view dump_func_regex,
+      mlir::OpPrintingFlags flags = mlir::OpPrintingFlags())
       : dump_dir_(dump_dir),
         dump_pass_re_(std::make_unique<RE2>(dump_pass_regex)),
-        dump_func_re_(std::make_unique<RE2>(dump_func_regex)) {}
+        dump_func_re_(std::make_unique<RE2>(dump_func_regex)),
+        op_printing_flags_(flags) {}
 
   DumpInstrumentation(const DumpInstrumentation& other) = delete;
   DumpInstrumentation& operator=(const DumpInstrumentation& other) = delete;
@@ -257,12 +262,13 @@ class DumpInstrumentation : public mlir::PassInstrumentation {
     file = std::make_unique<tsl::BufferedWritableFile>(std::move(file));
 
     WritableFileRawStream os(std::move(file));
-    op->print(os);
+    op->print(os, op_printing_flags_);
   }
 
   const std::string dump_dir_;
   const std::unique_ptr<RE2> dump_pass_re_;
   const std::unique_ptr<RE2> dump_func_re_;
+  mlir::OpPrintingFlags op_printing_flags_;
 
   // Counter used for pass name prefix to signify sequence
   int pass_counter_ = 0;
@@ -270,6 +276,40 @@ class DumpInstrumentation : public mlir::PassInstrumentation {
   bool printed_ = false;
 };
 
+double GetCurrentRssMb() {
+#if defined(__linux__)
+  std::ifstream statm("/proc/self/statm");
+  if (!statm.is_open()) return 0.0;
+  int64_t pages = 0;
+  int64_t rss_pages = 0;
+  if (statm >> pages >> rss_pages) {
+    int64_t page_size = sysconf(_SC_PAGESIZE);
+    return static_cast<double>(rss_pages * page_size) / (1024.0 * 1024.0);
+  }
+#endif
+  return 0.0;
+}
+
+class RssLoggingInstrumentation : public mlir::PassInstrumentation {
+ public:
+  void runBeforePass(mlir::Pass* pass, mlir::Operation* op) override {
+    pass_start_rss_[pass] = GetCurrentRssMb();
+  }
+
+  void runAfterPass(mlir::Pass* pass, mlir::Operation* op) override {
+    double end_rss = GetCurrentRssMb();
+    auto node = pass_start_rss_.extract(pass);
+    double start_rss = node.empty() ? 0.0 : node.mapped();
+    double delta = end_rss - start_rss;
+    LOG(INFO) << "[MLIR RSS] After pass '" << pass->getName().str() << "' on '"
+              << op->getName().getStringRef().str() << "': RSS = " << end_rss
+              << " MB (Delta: " << (delta >= 0 ? "+" : "") << delta << " MB)";
+  }
+
+ private:
+  absl::flat_hash_map<mlir::Pass*, double> pass_start_rss_;
+};
+
 std::function<bool(mlir::Pass*, mlir::Operation*)> CreatePrintIRFun(
     const std::string& pass_regex) {
   std::function<bool(mlir::Pass*, mlir::Operation*)> fun;
@@ -296,12 +336,20 @@ void InitPassManager(mlir::PassManager& pm,
   bool print_to_stdout =
       !options.print_ir_before().empty() || !options.print_ir_after().empty();
 
-  if (dump_to_dir || print_to_stdout) {
-    // Necessary for maintaining sequence of passes when dumping MLIR to files
-    // or stdout.
+  if (dump_to_dir || print_to_stdout || options.log_rss()) {
+    // Necessary for maintaining sequence of passes when dumping MLIR to files,
+    // stdout, or logging RSS memory.
     pm.getContext()->disableMultithreading();
   }
 
+  mlir::OpPrintingFlags opPrintingFlags = mlir::OpPrintingFlags();
+  if (options.has_elide_elementsattrs_if_larger()) {
+    opPrintingFlags.elideLargeElementsAttrs(
+        options.elide_elementsattrs_if_larger());
+    opPrintingFlags.elideLargeResourceString(
+        options.elide_elementsattrs_if_larger());
+  }
+
   if (dump_to_dir) {
     dump_dir = tsl::io::JoinPath(
         dump_dir, absl::FormatTime("%E4Y%m%d_%H%M%E6S", absl::Now(),
@@ -321,7 +369,8 @@ void InitPassManager(mlir::PassManager& pm,
     }
 
     pm.addInstrumentation(std::make_unique<DumpInstrumentation>(
-        dump_dir, options.ir_dump_pass_regex(), options.ir_dump_func_regex()));
+        dump_dir, options.ir_dump_pass_regex(), options.ir_dump_func_regex(),
+        opPrintingFlags));
   }
 
   if (print_to_stdout) {
@@ -331,13 +380,6 @@ void InitPassManager(mlir::PassManager& pm,
     std::function<bool(mlir::Pass*, mlir::Operation*)>
         should_print_ir_after_pass(CreatePrintIRFun(options.print_ir_after()));
 
-    mlir::OpPrintingFlags opPrintingFlags = mlir::OpPrintingFlags();
-
-    if (options.has_elide_elementsattrs_if_larger()) {
-      opPrintingFlags.elideLargeElementsAttrs(
-          options.elide_elementsattrs_if_larger());
-    }
-
     pm.enableIRPrinting(should_print_ir_before_pass, should_print_ir_after_pass,
                         options.print_ir_module_scope(),
                         /*printAfterOnlyOnChange=*/true,
@@ -345,6 +387,10 @@ void InitPassManager(mlir::PassManager& pm,
                         opPrintingFlags);
   }
 
+  if (options.log_rss()) {
+    pm.addInstrumentation(std::make_unique<RssLoggingInstrumentation>());
+  }
+
   // Enable pass timing. Note: MLIR expects `mlir::PassManager::enableTiming` to
   // be called after all instrumentations are added.
   if (options.enable_timing()) {
diff --git a/tensorflow/compiler/mlir/lite/debug/debug_options.proto b/tensorflow/compiler/mlir/lite/debug/debug_options.proto
index e45c1cd532a687..b085ea38f9aff9 100644
--- a/tensorflow/compiler/mlir/lite/debug/debug_options.proto
+++ b/tensorflow/compiler/mlir/lite/debug/debug_options.proto
@@ -19,7 +19,7 @@ package tensorflow.converter;
 
 // Additional parameters that control the debug behavior of the Converter.
 //
-// Next ID: 9
+// Next ID: 10
 message DebugOptions {
   // If not empty, dumps MLIR to the specified directory. The initial state of
   // the MLIR after import will be dumped at the beginning of each pass manager
@@ -58,4 +58,7 @@ message DebugOptions {
   // Elide ElementsAttrs with \"...\" that have more elements than the given
   // upper limit.
   optional int64 elide_elementsattrs_if_larger = 8;
+
+  // If true, log RSS memory usage (in MB) before and after each MLIR pass.
+  optional bool log_rss = 9 [default = false];
 }
diff --git a/tensorflow/compiler/mlir/lite/litert_mlir_reduce_main.cc b/tensorflow/compiler/mlir/lite/litert_mlir_reduce_main.cc
new file mode 100644
index 00000000000000..3cf43e4ed7fa36
--- /dev/null
+++ b/tensorflow/compiler/mlir/lite/litert_mlir_reduce_main.cc
@@ -0,0 +1,41 @@
+/* Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "llvm/Support/LogicalResult.h"
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/InitAllPasses.h"  // from @llvm-project
+#include "mlir/Tools/mlir-reduce/MlirReduceMain.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/init_mlir.h"
+#include "tensorflow/compiler/mlir/lite/register_lite_dialects.h"
+#include "tensorflow/compiler/mlir/lite/transforms/passes.h"
+#include "tensorflow/compiler/mlir/register_common_dialects.h"
+#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
+
+int main(int argc, char** argv) {
+  tensorflow::InitMlir y(&argc, &argv);
+
+  mlir::registerAllPasses();
+  mlir::registerTensorFlowPasses();
+  mlir::TFL::registerTensorFlowLitePasses();
+
+  mlir::DialectRegistry registry;
+  mlir::RegisterCommonToolingDialects(registry);
+  tflite::RegisterLiteToolingDialects(registry);
+
+  mlir::MLIRContext context;
+  context.appendDialectRegistry(registry);
+
+  return failed(mlir::mlirReduceMain(argc, argv, context));
+}
diff --git a/tensorflow/compiler/mlir/lite/python/BUILD b/tensorflow/compiler/mlir/lite/python/BUILD
index 44b2ef7416e28e..8411fbe7b0e56a 100644
--- a/tensorflow/compiler/mlir/lite/python/BUILD
+++ b/tensorflow/compiler/mlir/lite/python/BUILD
@@ -55,6 +55,7 @@ cc_library(
         "//tensorflow/compiler/mlir/lite:tf_to_tfl_flatbuffer",
         "//tensorflow/compiler/mlir/lite:types_proto_cc",
         "//tensorflow/compiler/mlir/lite/quantization/common/quantization_lib:quantization_config",
+        "//tensorflow/compiler/mlir/lite/quantization/ir:QuantOps",
         "//tensorflow/compiler/mlir/lite/tools/optimize:reduced_precision_metadata",
         "//tensorflow/compiler/mlir/quantization/tensorflow/python:py_function_lib",
         "//tensorflow/compiler/mlir/tensorflow:tensorflow_ops",
@@ -79,11 +80,41 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "conversion_failure_reporter",
+    srcs = [
+        "conversion_failure_reporter.cc",
+        "pass_debug_instrumentation.cc",
+    ],
+    hdrs = [
+        "conversion_failure_reporter.h",
+        "pass_debug_instrumentation.h",
+    ],
+    deps = [
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:str_format",
+        "@com_google_absl//absl/time",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:BytecodeWriter",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+        "@xla//xla/mlir/utils:error_util",
+        "@xla//xla/tsl/platform:env",
+    ],
+)
+
 cc_library(
     name = "stablehlo_tfl_pipeline",
-    srcs = ["stablehlo_tfl_pipeline.cc"],
-    hdrs = ["stablehlo_tfl_pipeline.h"],
+    srcs = [
+        "stablehlo_tfl_pipeline.cc",
+    ],
+    hdrs = [
+        "stablehlo_tfl_pipeline.h",
+    ],
     deps = [
+        ":conversion_failure_reporter",
         "//tensorflow/compiler/mlir/lite:common",
         "//tensorflow/compiler/mlir/lite:converter_flags_proto_cc",
         "//tensorflow/compiler/mlir/lite:flatbuffer_export",
@@ -97,8 +128,11 @@ cc_library(
         "//tensorflow/compiler/mlir/lite/stablehlo:prepare_hlo",
         "//tensorflow/compiler/mlir/lite/stablehlo:tfl_legalize_hlo",
         "//tensorflow/compiler/mlir/lite/stablehlo:unfold_splat_constant_pass",
+        "//tensorflow/compiler/mlir/tensorflow:error_util",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:str_format",
+        "@com_google_absl//absl/time",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:BytecodeWriter",
         "@llvm-project//mlir:FuncDialect",
@@ -112,9 +146,25 @@ cc_library(
         "@stablehlo//:stablehlo_ops",
         "@stablehlo//:stablehlo_passes",
         "@stablehlo//:vhlo_ops",
+        "@xla//xla/mlir/utils:error_util",
         "@xla//xla/mlir_hlo",
         "@xla//xla/mlir_hlo:mhlo_passes",
         "@xla//xla/mlir_hlo:stablehlo_extension_passes",
+        "@xla//xla/tsl/platform:env",
+    ],
+)
+
+tf_cc_test(
+    name = "conversion_failure_reporter_test",
+    srcs = ["conversion_failure_reporter_test.cc"],
+    deps = [
+        ":conversion_failure_reporter",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+        "@com_google_googletest//:gtest_main",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Parser",
+        "@xla//xla/tsl/platform:env",
     ],
 )
 
@@ -177,6 +227,7 @@ cc_library(
     hdrs = ["slim_model_importer.h"],
     deps = [
         "//tensorflow/compiler/mlir/lite/stablehlo:drop_shape_assertions",
+        "//tensorflow/compiler/mlir/lite/stablehlo:legalize_vhlo_quant_custom_calls",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
@@ -186,6 +237,7 @@ cc_library(
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:Parser",
         "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:ReconcileUnrealizedCasts",
         "@llvm-project//mlir:Support",
         "@stablehlo//:stablehlo_ops",
         "@stablehlo//:stablehlo_passes",
diff --git a/tensorflow/compiler/mlir/lite/python/conversion_failure_reporter.cc b/tensorflow/compiler/mlir/lite/python/conversion_failure_reporter.cc
new file mode 100644
index 00000000000000..38c2de342ef318
--- /dev/null
+++ b/tensorflow/compiler/mlir/lite/python/conversion_failure_reporter.cc
@@ -0,0 +1,278 @@
+/* Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/lite/python/conversion_failure_reporter.h"
+
+#if defined(_WIN32)
+#include <process.h>
+#else
+#include <unistd.h>
+#endif
+
+#include <cstddef>
+#include <cstdint>
+#include <fstream>
+#include <ios>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/strings/ascii.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/str_format.h"
+#include "absl/strings/str_split.h"
+#include "absl/strings/string_view.h"
+#include "absl/time/clock.h"
+#include "absl/time/time.h"
+#include "llvm/Support/FormatVariadic.h"
+#include "llvm/Support/JSON.h"
+#include "llvm/Support/raw_os_ostream.h"
+#include "llvm/Support/raw_ostream.h"
+#include "mlir/Bytecode/BytecodeWriter.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/OperationSupport.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "xla/tsl/platform/env.h"
+
+namespace mlir::TFL {
+namespace {
+
+bool IsCaretLine(absl::string_view line) {
+  absl::string_view trimmed = absl::StripAsciiWhitespace(line);
+  return trimmed == "^";
+}
+
+}  // namespace
+
+FailureReport ConversionFailureReporter::ParseDiagnostic(
+    absl::string_view raw_error, absl::string_view stage,
+    absl::string_view status_code, absl::string_view failing_pass,
+    absl::string_view failing_pass_arg, absl::string_view failing_function) {
+  FailureReport report;
+  report.stage = std::string(stage);
+  report.status_code = std::string(status_code);
+  report.failing_pass = std::string(failing_pass);
+  report.failing_pass_arg = std::string(failing_pass_arg);
+  report.failing_function = std::string(failing_function);
+  report.raw_error = std::string(raw_error);
+
+  std::vector<std::string> lines = absl::StrSplit(raw_error, '\n');
+
+  enum State { kStart, kPrimarySnippet, kCallStackSnippet, kOpContextSnippet };
+  State state = kStart;
+
+  for (size_t i = 0; i < lines.size(); ++i) {
+    const std::string& line = lines[i];
+
+    // Check for primary error line: "... error: ..."
+    auto err_pos = line.find(": error: ");
+    if (err_pos != std::string::npos) {
+      if (report.primary_error.message.empty()) {
+        std::string loc = line.substr(0, err_pos);
+        auto file_pos = loc.find("third_party/");
+        if (file_pos != std::string::npos) {
+          loc = loc.substr(file_pos);
+        }
+        report.primary_error.location = loc;
+        report.primary_error.severity = "error";
+        report.primary_error.message = line.substr(err_pos + 9);
+        report.primary_error.code_snippet.clear();
+        state = kPrimarySnippet;
+      } else {
+        // Subsequent errors should not mix into the first primary error's
+        // snippet.
+        state = kStart;
+      }
+      continue;
+    }
+
+    // Check for call stack note: "... note: called from"
+    auto call_pos = line.find(": note: called from");
+    if (call_pos != std::string::npos) {
+      FailureReport::CallStackFrame frame;
+      std::string loc = line.substr(0, call_pos);
+      auto file_pos = loc.find("third_party/");
+      if (file_pos != std::string::npos) {
+        loc = loc.substr(file_pos);
+      }
+      frame.location = loc;
+      report.call_stack.push_back(frame);
+      state = kCallStackSnippet;
+      continue;
+    }
+
+    // Check for operation context note: "... note: see current operation: "
+    auto op_pos = line.find(": note: see current operation: ");
+    if (op_pos != std::string::npos) {
+      std::string loc = line.substr(0, op_pos);
+      auto file_pos = loc.find("third_party/");
+      if (file_pos != std::string::npos) {
+        loc = loc.substr(file_pos);
+      }
+      report.operation_context.location = loc;
+      report.operation_context.operation = line.substr(op_pos + 31);
+      state = kOpContextSnippet;
+      continue;
+    }
+
+    // Otherwise, code snippet or continuation line
+    absl::string_view trimmed_line = absl::StripAsciiWhitespace(line);
+    if (!trimmed_line.empty() && !IsCaretLine(trimmed_line)) {
+      if (state == kPrimarySnippet) {
+        if (!report.primary_error.code_snippet.empty()) {
+          absl::StrAppend(&report.primary_error.code_snippet, "\n");
+        }
+        absl::StrAppend(&report.primary_error.code_snippet, trimmed_line);
+      } else if (state == kCallStackSnippet && !report.call_stack.empty()) {
+        if (!report.call_stack.back().code_snippet.empty()) {
+          absl::StrAppend(&report.call_stack.back().code_snippet, "\n");
+        }
+        absl::StrAppend(&report.call_stack.back().code_snippet, trimmed_line);
+      } else if (state == kOpContextSnippet) {
+        if (!report.operation_context.operation.empty()) {
+          absl::StrAppend(&report.operation_context.operation, "\n");
+        }
+        absl::StrAppend(&report.operation_context.operation, trimmed_line);
+      }
+    }
+  }
+
+  return report;
+}
+
+std::string ConversionFailureReporter::GetOrCreateDebugDir(
+    absl::string_view working_dir) {
+  if (!working_dir.empty()) {
+    return std::string(working_dir);
+  }
+#if defined(_WIN32)
+  int pid = _getpid();
+#else
+  int pid = getpid();
+#endif
+  return absl::StrFormat(
+      "/tmp/litert_conv_%s_%d",
+      absl::FormatTime("%E4Y%m%d_%H%M%E6S", absl::Now(), absl::LocalTimeZone()),
+      pid);
+}
+
+void ConversionFailureReporter::WriteFailureJson(
+    absl::string_view working_dir, mlir::ModuleOp module,
+    absl::string_view raw_error, absl::string_view stage,
+    absl::string_view status_code, absl::string_view failing_pass,
+    absl::string_view failing_pass_arg, bool write_module_artifacts,
+    absl::string_view failing_function, int64_t elide_elements_larger_than,
+    int64_t elide_resource_strings_larger_than) {
+  std::string dir = GetOrCreateDebugDir(working_dir);
+  if (!tsl::Env::Default()->RecursivelyCreateDir(dir).ok()) return;
+
+  FailureReport report =
+      ParseDiagnostic(raw_error, stage, status_code, failing_pass,
+                      failing_pass_arg, failing_function);
+
+  // Dump elided MLIR text module if module is valid and artifacts requested
+  if (module && write_module_artifacts) {
+    std::string elided_path = absl::StrCat(dir, "/before_failure_elided.mlir");
+    std::ofstream elided_file(elided_path);
+    if (elided_file.is_open()) {
+      llvm::raw_os_ostream elided_os(elided_file);
+      mlir::OpPrintingFlags flags;
+      flags.elideLargeElementsAttrs(elide_elements_larger_than);
+      flags.elideLargeResourceString(elide_resource_strings_larger_than);
+      module.print(elided_os, flags);
+      elided_os.flush();
+      elided_file.flush();
+      report.elided_mlir_file = elided_path;
+    }
+
+    // Dump non-elided binary MLIR bytecode module
+    std::string bc_path = absl::StrCat(dir, "/module_bytecode.mlirbc");
+    std::ofstream bc_file(bc_path, std::ios::binary);
+    if (bc_file.is_open()) {
+      llvm::raw_os_ostream bc_os(bc_file);
+      if (mlir::succeeded(mlir::writeBytecodeToFile(module, bc_os))) {
+        bc_os.flush();
+        bc_file.flush();
+        report.bytecode_file = bc_path;
+      }
+    }
+  }
+
+  std::string file_path = absl::StrCat(dir, "/failure.json");
+
+  llvm::json::Object root;
+  root["stage"] = report.stage;
+  root["status_code"] = report.status_code;
+
+  if (!report.failing_pass.empty()) {
+    root["failing_pass"] = report.failing_pass;
+  }
+  if (!report.failing_pass_arg.empty()) {
+    root["failing_pass_arg"] = report.failing_pass_arg;
+  }
+  if (!report.failing_function.empty()) {
+    root["failing_function"] = report.failing_function;
+  }
+
+  if (!report.elided_mlir_file.empty()) {
+    root["elided_mlir_file"] = report.elided_mlir_file;
+  }
+  if (!report.bytecode_file.empty()) {
+    root["bytecode_file"] = report.bytecode_file;
+  }
+
+  if (!report.primary_error.message.empty()) {
+    llvm::json::Object prim_err;
+    prim_err["location"] = report.primary_error.location;
+    prim_err["severity"] = report.primary_error.severity;
+    prim_err["message"] = report.primary_error.message;
+    if (!report.primary_error.code_snippet.empty()) {
+      prim_err["code_snippet"] = report.primary_error.code_snippet;
+    }
+    root["primary_error"] = std::move(prim_err);
+  }
+
+  if (!report.call_stack.empty()) {
+    llvm::json::Array stack_arr;
+    for (const auto& frame : report.call_stack) {
+      llvm::json::Object frame_obj;
+      frame_obj["location"] = frame.location;
+      if (!frame.code_snippet.empty()) {
+        frame_obj["code_snippet"] = frame.code_snippet;
+      }
+      stack_arr.push_back(std::move(frame_obj));
+    }
+    root["call_stack"] = std::move(stack_arr);
+  }
+
+  if (!report.operation_context.operation.empty()) {
+    llvm::json::Object op_ctx;
+    op_ctx["location"] = report.operation_context.location;
+    op_ctx["operation"] = report.operation_context.operation;
+    root["operation_context"] = std::move(op_ctx);
+  }
+
+  root["raw_error"] = report.raw_error;
+
+  std::ofstream json_file(file_path);
+  if (json_file.is_open()) {
+    llvm::raw_os_ostream os(json_file);
+    os << llvm::formatv("{0:2}\n", llvm::json::Value(std::move(root)));
+    os.flush();
+    json_file.flush();
+  }
+}
+
+}  // namespace mlir::TFL
diff --git a/tensorflow/compiler/mlir/lite/python/conversion_failure_reporter.h b/tensorflow/compiler/mlir/lite/python/conversion_failure_reporter.h
new file mode 100644
index 00000000000000..0f49308b488d94
--- /dev/null
+++ b/tensorflow/compiler/mlir/lite/python/conversion_failure_reporter.h
@@ -0,0 +1,87 @@
+/* Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_LITE_PYTHON_CONVERSION_FAILURE_REPORTER_H_
+#define TENSORFLOW_COMPILER_MLIR_LITE_PYTHON_CONVERSION_FAILURE_REPORTER_H_
+
+#include <cstdint>
+#include <string>
+#include <vector>
+
+#include "absl/strings/string_view.h"
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+
+namespace mlir::TFL {
+
+struct FailureReport {
+  std::string stage;
+  std::string status_code;
+  std::string failing_pass;
+  std::string failing_pass_arg;
+  std::string failing_function;
+  std::string raw_error;
+  std::string elided_mlir_file;
+  std::string bytecode_file;
+
+  struct PrimaryError {
+    std::string location;
+    std::string severity;
+    std::string message;
+    std::string code_snippet;
+  } primary_error;
+
+  struct CallStackFrame {
+    std::string location;
+    std::string code_snippet;
+  };
+  std::vector<CallStackFrame> call_stack;
+
+  struct OperationContext {
+    std::string location;
+    std::string operation;
+  } operation_context;
+};
+
+class ConversionFailureReporter {
+ public:
+  // Parses a raw MLIR error diagnostic message into structured fields.
+  static FailureReport ParseDiagnostic(absl::string_view raw_error,
+                                       absl::string_view stage,
+                                       absl::string_view status_code,
+                                       absl::string_view failing_pass = "",
+                                       absl::string_view failing_pass_arg = "",
+                                       absl::string_view failing_function = "");
+
+  // Returns working_dir if non-empty, or generates a unique directory name
+  // under /tmp based on timestamp and process ID.
+  static std::string GetOrCreateDebugDir(absl::string_view working_dir);
+
+  // Writes the structured FailureReport, elided MLIR, and bytecode module.
+  static void WriteFailureJson(absl::string_view working_dir,
+                               mlir::ModuleOp module,
+                               absl::string_view raw_error,
+                               absl::string_view stage,
+                               absl::string_view status_code,
+                               absl::string_view failing_pass = "",
+                               absl::string_view failing_pass_arg = "",
+                               bool write_module_artifacts = true,
+                               absl::string_view failing_function = "",
+                               int64_t elide_elements_larger_than = 8,
+                               int64_t elide_resource_strings_larger_than = 64);
+};
+
+}  // namespace mlir::TFL
+
+#endif  // TENSORFLOW_COMPILER_MLIR_LITE_PYTHON_CONVERSION_FAILURE_REPORTER_H_
diff --git a/tensorflow/compiler/mlir/lite/python/conversion_failure_reporter_test.cc b/tensorflow/compiler/mlir/lite/python/conversion_failure_reporter_test.cc
new file mode 100644
index 00000000000000..7e8531ceddcb4a
--- /dev/null
+++ b/tensorflow/compiler/mlir/lite/python/conversion_failure_reporter_test.cc
@@ -0,0 +1,153 @@
+/* Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/lite/python/conversion_failure_reporter.h"
+
+#include <string>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "mlir/Parser/Parser.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/python/pass_debug_instrumentation.h"
+#include "xla/tsl/platform/env.h"
+
+namespace mlir::TFL {
+namespace {
+
+using ::testing::Eq;
+using ::testing::HasSubstr;
+using ::testing::IsEmpty;
+using ::testing::SizeIs;
+
+TEST(ConversionFailureReporterTest, GetOrCreateDebugDirCustom) {
+  std::string dir =
+      ConversionFailureReporter::GetOrCreateDebugDir("/my/custom/dir");
+  EXPECT_EQ(dir, "/my/custom/dir");
+}
+
+TEST(ConversionFailureReporterTest, GetOrCreateDebugDirDefaultDynamic) {
+  std::string dir = ConversionFailureReporter::GetOrCreateDebugDir("");
+  EXPECT_THAT(dir, HasSubstr("/tmp/litert_conv_"));
+}
+
+TEST(ConversionFailureReporterTest, ParseDiagnosticFull) {
+  const std::string raw_error = R"(
+third_party/models/test.py:42:1: error: custom error message
+  %val = "tf.UnsupportedOp"(%0)
+  ^
+third_party/models/caller.py:10:5: note: called from
+  return helper_function(x)
+third_party/models/graph.mlir:15:3: note: see current operation: "tfl.custom"(%arg0) {
+  attr1 = 42 : i32,
+  attr2 = "value"
+}
+)";
+
+  FailureReport report = ConversionFailureReporter::ParseDiagnostic(
+      raw_error, "PassPipeline", "INVALID_ARGUMENT", "TestPass",
+      "--test-pass-flag", "main_func");
+
+  EXPECT_EQ(report.stage, "PassPipeline");
+  EXPECT_EQ(report.status_code, "INVALID_ARGUMENT");
+  EXPECT_EQ(report.failing_pass, "TestPass");
+  EXPECT_EQ(report.failing_pass_arg, "--test-pass-flag");
+  EXPECT_EQ(report.failing_function, "main_func");
+
+  // Primary error
+  EXPECT_EQ(report.primary_error.location, "third_party/models/test.py:42:1");
+  EXPECT_EQ(report.primary_error.severity, "error");
+  EXPECT_EQ(report.primary_error.message, "custom error message");
+  EXPECT_EQ(report.primary_error.code_snippet,
+            "%val = \"tf.UnsupportedOp\"(%0)");
+
+  // Call stack
+  ASSERT_THAT(report.call_stack, SizeIs(1));
+  EXPECT_EQ(report.call_stack[0].location, "third_party/models/caller.py:10:5");
+  EXPECT_EQ(report.call_stack[0].code_snippet, "return helper_function(x)");
+
+  // Operation context with multi-line continuation
+  EXPECT_EQ(report.operation_context.location,
+            "third_party/models/graph.mlir:15:3");
+  EXPECT_THAT(report.operation_context.operation,
+              HasSubstr("\"tfl.custom\"(%arg0) {"));
+  EXPECT_THAT(report.operation_context.operation,
+              HasSubstr("attr1 = 42 : i32,"));
+  EXPECT_THAT(report.operation_context.operation,
+              HasSubstr("attr2 = \"value\""));
+}
+
+TEST(ConversionFailureReporterTest, WriteFailureJsonCreatesArtifacts) {
+  std::string test_dir;
+  ASSERT_TRUE(tsl::Env::Default()->LocalTempFilename(&test_dir));
+  ASSERT_TRUE(tsl::Env::Default()->RecursivelyCreateDir(test_dir).ok());
+
+  MLIRContext context;
+  OwningOpRef<ModuleOp> module =
+      parseSourceString<ModuleOp>("module {}", &context);
+  ASSERT_TRUE(module);
+
+  ConversionFailureReporter::WriteFailureJson(
+      test_dir, *module, "loc:1:1: error: failed pass", "TestStage",
+      "INVALID_ARGUMENT", "FailingPass", "--pass-arg",
+      /*write_module_artifacts=*/true, "main");
+
+  std::string failure_json_path = absl::StrCat(test_dir, "/failure.json");
+  std::string elided_mlir_path =
+      absl::StrCat(test_dir, "/before_failure_elided.mlir");
+  std::string bytecode_path = absl::StrCat(test_dir, "/module_bytecode.mlirbc");
+
+  EXPECT_TRUE(tsl::Env::Default()->FileExists(failure_json_path).ok());
+  EXPECT_TRUE(tsl::Env::Default()->FileExists(elided_mlir_path).ok());
+  EXPECT_TRUE(tsl::Env::Default()->FileExists(bytecode_path).ok());
+
+  std::string json_contents;
+  ASSERT_TRUE(tsl::ReadFileToString(tsl::Env::Default(), failure_json_path,
+                                    &json_contents)
+                  .ok());
+  EXPECT_THAT(json_contents, HasSubstr("\"stage\": \"TestStage\""));
+  EXPECT_THAT(json_contents, HasSubstr("\"failing_pass\": \"FailingPass\""));
+  EXPECT_THAT(json_contents, HasSubstr("\"failing_function\": \"main\""));
+}
+
+TEST(PipelineFailureCoordinatorTest, ReportSerializationFailure) {
+  std::string test_dir;
+  ASSERT_TRUE(tsl::Env::Default()->LocalTempFilename(&test_dir));
+  ASSERT_TRUE(tsl::Env::Default()->RecursivelyCreateDir(test_dir).ok());
+
+  MLIRContext context;
+  OwningOpRef<ModuleOp> module =
+      parseSourceString<ModuleOp>("module {}", &context);
+  ASSERT_TRUE(module);
+
+  PipelineFailureCoordinator coordinator(test_dir, /*enable_debug=*/true);
+  absl::Status status = coordinator.ReportSerializationFailure(
+      *module, absl::InternalError("internal serialization error"),
+      "diag details");
+
+  EXPECT_TRUE(absl::IsInvalidArgument(status));
+  EXPECT_THAT(status.message(),
+              HasSubstr("Failed to serialize to FlatBuffer: diag details"));
+
+  std::string failure_json_path = absl::StrCat(test_dir, "/failure.json");
+  EXPECT_TRUE(tsl::Env::Default()->FileExists(failure_json_path).ok());
+}
+
+}  // namespace
+}  // namespace mlir::TFL
diff --git a/tensorflow/compiler/mlir/lite/python/pass_debug_instrumentation.cc b/tensorflow/compiler/mlir/lite/python/pass_debug_instrumentation.cc
new file mode 100644
index 00000000000000..45c7ce8eaa6b57
--- /dev/null
+++ b/tensorflow/compiler/mlir/lite/python/pass_debug_instrumentation.cc
@@ -0,0 +1,222 @@
+/* Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/lite/python/pass_debug_instrumentation.h"
+
+#include <cstdint>
+#include <fstream>
+#include <ios>
+#include <memory>
+#include <string>
+
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/str_format.h"
+#include "absl/strings/string_view.h"
+#include "llvm/Support/Regex.h"
+#include "llvm/Support/raw_os_ostream.h"
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/OperationSupport.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Pass/PassInstrumentation.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/python/conversion_failure_reporter.h"
+#include "xla/tsl/platform/env.h"
+
+namespace mlir::TFL {
+
+PassDebugInstrumentation::PassDebugInstrumentation(
+    std::string* name, std::string* arg, std::string* failing_func,
+    mlir::OwningOpRef<mlir::ModuleOp>* clean_module, bool debug,
+    absl::string_view dir, absl::string_view print_before,
+    absl::string_view print_after, int64_t elide_elements_larger_than,
+    int64_t elide_resource_strings_larger_than)
+    : name_(name),
+      arg_(arg),
+      failing_func_(failing_func),
+      clean_module_(clean_module),
+      enable_debug_(debug),
+      debug_dir_(dir),
+      elide_elements_larger_than_(elide_elements_larger_than),
+      elide_resource_strings_larger_than_(elide_resource_strings_larger_than) {
+  if (!print_before.empty()) {
+    print_before_regex_ = std::make_unique<llvm::Regex>(print_before);
+    std::string error;
+    if (!print_before_regex_->isValid(error)) {
+      print_before_regex_.reset();
+    }
+  }
+  if (!print_after.empty()) {
+    print_after_regex_ = std::make_unique<llvm::Regex>(print_after);
+    std::string error;
+    if (!print_after_regex_->isValid(error)) {
+      print_after_regex_.reset();
+    }
+  }
+}
+
+void PassDebugInstrumentation::DumpIrToFile(mlir::Pass* pass,
+                                            mlir::Operation* op,
+                                            llvm::StringRef suffix) {
+  if (!op) return;
+  std::string dumps_dir = absl::StrCat(debug_dir_, "/ir_dumps");
+  if (!tsl::Env::Default()->RecursivelyCreateDir(dumps_dir).ok()) return;
+  std::string func_name;
+  if (auto sym_attr = op->getAttrOfType<mlir::StringAttr>("sym_name")) {
+    func_name =
+        absl::StrCat("_", absl::string_view(sym_attr.getValue().data(),
+                                            sym_attr.getValue().size()));
+  }
+  std::string filename = absl::StrFormat(
+      "%s/%04d_%s%s_%s.mlir", dumps_dir, step_counter_,
+      pass ? pass->getName().str() : "unknown", func_name, suffix.str());
+
+  std::ofstream ir_file(filename);
+  if (ir_file.is_open()) {
+    llvm::raw_os_ostream os(ir_file);
+    mlir::OpPrintingFlags flags;
+    flags.elideLargeElementsAttrs(elide_elements_larger_than_);
+    flags.elideLargeResourceString(elide_resource_strings_larger_than_);
+    op->print(os, flags);
+    os.flush();
+    ir_file.flush();
+  }
+}
+
+bool PassDebugInstrumentation::MatchPassRegex(llvm::StringRef pass_name,
+                                              const llvm::Regex* regex) {
+  if (!regex) return false;
+  return regex->match(pass_name);
+}
+
+void PassDebugInstrumentation::runBeforePass(mlir::Pass* pass,
+                                             mlir::Operation* op) {
+  if (pass && pass->getName() == "mlir::detail::OpToOpPassAdaptor") {
+    return;
+  }
+  step_counter_++;
+  std::string func_name;
+  if (op) {
+    if (auto sym_attr = op->getAttrOfType<mlir::StringAttr>("sym_name")) {
+      func_name = sym_attr.getValue().str();
+    } else {
+      func_name = op->getName().getStringRef().str();
+    }
+  }
+  if (enable_debug_) {
+    std::string trace_file =
+        absl::StrCat(debug_dir_, "/pass_execution_trace.log");
+    std::ofstream trace_os(trace_file, std::ios::app);
+    if (trace_os.is_open()) {
+      trace_os << "[" << step_counter_
+               << "] Pass: " << (pass ? pass->getName().str() : "unknown")
+               << " | Op: " << func_name << "\n";
+    }
+
+    if (clean_module_ && op) {
+      if (auto mod = mlir::dyn_cast<mlir::ModuleOp>(op)) {
+        *clean_module_ = mod.clone();
+      } else if (auto mod = op->getParentOfType<mlir::ModuleOp>()) {
+        *clean_module_ = mod.clone();
+      }
+    }
+    if (pass && MatchPassRegex(pass->getName(), print_before_regex_.get())) {
+      DumpIrToFile(pass, op, "before");
+    }
+  }
+}
+
+void PassDebugInstrumentation::runAfterPass(mlir::Pass* pass,
+                                            mlir::Operation* op) {
+  if (pass && pass->getName() == "mlir::detail::OpToOpPassAdaptor") {
+    return;
+  }
+  if (enable_debug_ && pass &&
+      MatchPassRegex(pass->getName(), print_after_regex_.get())) {
+    DumpIrToFile(pass, op, "after");
+  }
+}
+
+void PassDebugInstrumentation::runAfterPassFailed(mlir::Pass* pass,
+                                                  mlir::Operation* op) {
+  if (pass && pass->getName() == "mlir::detail::OpToOpPassAdaptor") {
+    return;
+  }
+  if (pass && name_ && name_->empty()) {
+    *name_ = std::string(pass->getName());
+    if (arg_) *arg_ = std::string(pass->getArgument());
+  }
+  std::string func_name;
+  if (op) {
+    if (auto sym_attr = op->getAttrOfType<mlir::StringAttr>("sym_name")) {
+      func_name = sym_attr.getValue().str();
+    } else {
+      func_name = op->getName().getStringRef().str();
+    }
+  }
+  if (!func_name.empty() && failing_func_ && failing_func_->empty()) {
+    *failing_func_ = func_name;
+  }
+}
+
+std::unique_ptr<mlir::PassInstrumentation>
+PipelineFailureCoordinator::CreateInstrumentation(
+    absl::string_view print_before_pattern,
+    absl::string_view print_after_pattern) {
+  return std::make_unique<PassDebugInstrumentation>(
+      &failing_pass_name_, &failing_pass_arg_, &failing_function_name_,
+      &pre_pass_clean_module_, enable_debug_, debug_dir_, print_before_pattern,
+      print_after_pattern, elide_elements_larger_than_,
+      elide_resource_strings_larger_than_);
+}
+
+absl::Status PipelineFailureCoordinator::ReportFailure(
+    mlir::ModuleOp fallback_module, const absl::Status& pass_status) const {
+  std::string err_msg =
+      pass_status.ok() ? "StableHLO to TFLite pipeline failed."
+                       : absl::StrCat("StableHLO to TFLite pipeline failed: ",
+                                      pass_status.message());
+  std::string pass_arg_flag =
+      failing_pass_arg_.empty() ? "" : absl::StrCat("--", failing_pass_arg_);
+  mlir::ModuleOp dump_module =
+      pre_pass_clean_module_ ? pre_pass_clean_module_.get() : fallback_module;
+  ConversionFailureReporter::WriteFailureJson(
+      debug_dir_, dump_module, err_msg, "StableHLO_to_TFLite_Pass_Pipeline",
+      absl::StatusCodeToString(pass_status.code()), failing_pass_name_,
+      pass_arg_flag, /*write_module_artifacts=*/enable_debug_,
+      failing_function_name_, elide_elements_larger_than_,
+      elide_resource_strings_larger_than_);
+  return absl::InvalidArgumentError(err_msg);
+}
+
+absl::Status PipelineFailureCoordinator::ReportSerializationFailure(
+    mlir::ModuleOp module, const absl::Status& status,
+    absl::string_view diag_errors) const {
+  std::string detail(diag_errors.empty() ? status.message() : diag_errors);
+  std::string err_msg =
+      absl::StrCat("Failed to serialize to FlatBuffer: ", detail);
+  ConversionFailureReporter::WriteFailureJson(
+      debug_dir_, module, err_msg, "FlatBuffer_Serialization",
+      absl::StatusCodeToString(status.code()),
+      /*failing_pass=*/"", /*failing_pass_arg=*/"",
+      /*write_module_artifacts=*/enable_debug_,
+      /*failing_function=*/"", elide_elements_larger_than_,
+      elide_resource_strings_larger_than_);
+  return absl::InvalidArgumentError(err_msg);
+}
+
+}  // namespace mlir::TFL
diff --git a/tensorflow/compiler/mlir/lite/python/pass_debug_instrumentation.h b/tensorflow/compiler/mlir/lite/python/pass_debug_instrumentation.h
new file mode 100644
index 00000000000000..85665bf05774e4
--- /dev/null
+++ b/tensorflow/compiler/mlir/lite/python/pass_debug_instrumentation.h
@@ -0,0 +1,124 @@
+/* Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_LITE_PYTHON_PASS_DEBUG_INSTRUMENTATION_H_
+#define TENSORFLOW_COMPILER_MLIR_LITE_PYTHON_PASS_DEBUG_INSTRUMENTATION_H_
+
+#include <cstdint>
+#include <memory>
+#include <string>
+
+#include "absl/status/status.h"
+#include "absl/strings/string_view.h"
+#include "llvm/Support/Regex.h"
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/Diagnostics.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Pass/PassInstrumentation.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "xla/mlir/utils/error_util.h"
+
+namespace mlir::TFL {
+
+class PassDebugInstrumentation : public mlir::PassInstrumentation {
+ public:
+  PassDebugInstrumentation(std::string* name, std::string* arg,
+                           std::string* failing_func,
+                           mlir::OwningOpRef<mlir::ModuleOp>* clean_module,
+                           bool enable_debug, absl::string_view debug_dir,
+                           absl::string_view print_before,
+                           absl::string_view print_after,
+                           int64_t elide_elements_larger_than = 8,
+                           int64_t elide_resource_strings_larger_than = 64);
+
+  void runBeforePass(mlir::Pass* pass, mlir::Operation* op) override;
+  void runAfterPass(mlir::Pass* pass, mlir::Operation* op) override;
+  void runAfterPassFailed(mlir::Pass* pass, mlir::Operation* op) override;
+
+ private:
+  void DumpIrToFile(mlir::Pass* pass, mlir::Operation* op,
+                    llvm::StringRef suffix);
+  bool MatchPassRegex(llvm::StringRef pass_name, const llvm::Regex* regex);
+
+  std::string* name_;
+  std::string* arg_;
+  std::string* failing_func_;
+  mlir::OwningOpRef<mlir::ModuleOp>* clean_module_;
+  bool enable_debug_;
+  std::string debug_dir_;
+  std::unique_ptr<llvm::Regex> print_before_regex_;
+  std::unique_ptr<llvm::Regex> print_after_regex_;
+  int64_t elide_elements_larger_than_ = 8;
+  int64_t elide_resource_strings_larger_than_ = 64;
+  int step_counter_ = 0;
+};
+
+class SerializationDiagHandler : public mlir::BaseScopedDiagnosticHandler {
+ public:
+  explicit SerializationDiagHandler(mlir::MLIRContext* ctx, std::string* out)
+      : BaseScopedDiagnosticHandler(ctx), out_(out) {
+    setHandler([this](mlir::Diagnostic& diag) {
+      if (diag.getSeverity() == mlir::DiagnosticSeverity::Error) {
+        if (!out_->empty()) *out_ += '\n';
+        *out_ += diag.str();
+      }
+      return mlir::failure();
+    });
+  }
+
+ private:
+  std::string* out_;
+};
+
+// Coordinating class that encapsulates pass failure diagnostic variables and
+// coordinates with PassDebugInstrumentation and ConversionFailureReporter.
+class PipelineFailureCoordinator {
+ public:
+  PipelineFailureCoordinator(const std::string& debug_dir, bool enable_debug,
+                             int64_t elide_elements_larger_than = 8,
+                             int64_t elide_resource_strings_larger_than = 64)
+      : debug_dir_(debug_dir),
+        enable_debug_(enable_debug),
+        elide_elements_larger_than_(elide_elements_larger_than),
+        elide_resource_strings_larger_than_(
+            elide_resource_strings_larger_than) {}
+
+  std::unique_ptr<mlir::PassInstrumentation> CreateInstrumentation(
+      absl::string_view print_before_pattern,
+      absl::string_view print_after_pattern);
+
+  absl::Status ReportFailure(mlir::ModuleOp fallback_module,
+                             const absl::Status& pass_status) const;
+
+  absl::Status ReportSerializationFailure(mlir::ModuleOp module,
+                                          const absl::Status& status,
+                                          absl::string_view diag_errors) const;
+
+ private:
+  std::string debug_dir_;
+  bool enable_debug_;
+  int64_t elide_elements_larger_than_ = 8;
+  int64_t elide_resource_strings_larger_than_ = 64;
+  std::string failing_pass_name_;
+  std::string failing_pass_arg_;
+  std::string failing_function_name_;
+  mlir::OwningOpRef<mlir::ModuleOp> pre_pass_clean_module_;
+};
+
+}  // namespace mlir::TFL
+
+#endif  // TENSORFLOW_COMPILER_MLIR_LITE_PYTHON_PASS_DEBUG_INSTRUMENTATION_H_
diff --git a/tensorflow/compiler/mlir/lite/python/stablehlo_tfl_pipeline.cc b/tensorflow/compiler/mlir/lite/python/stablehlo_tfl_pipeline.cc
index 59504630a1db05..0b64e690b49e38 100644
--- a/tensorflow/compiler/mlir/lite/python/stablehlo_tfl_pipeline.cc
+++ b/tensorflow/compiler/mlir/lite/python/stablehlo_tfl_pipeline.cc
@@ -15,12 +15,16 @@ limitations under the License.
 
 #include "tensorflow/compiler/mlir/lite/python/stablehlo_tfl_pipeline.h"
 
-#include <unistd.h>
-
+#include <cstdint>
+#include <fstream>
 #include <memory>
+#include <string>
+#include <utility>
 
 #include "absl/status/status.h"
 #include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "llvm/Support/raw_os_ostream.h"
 #include "llvm/Support/raw_ostream.h"
 #include "mlir/Bytecode/BytecodeWriter.h"  // from @llvm-project
 #include "mlir/Conversion/ReconcileUnrealizedCasts/ReconcileUnrealizedCasts.h"  // from @llvm-project
@@ -32,6 +36,7 @@ limitations under the License.
 #include "mlir/Pass/PassInstrumentation.h"  // from @llvm-project
 #include "mlir/Pass/PassManager.h"  // from @llvm-project
 #include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/Timing.h"  // from @llvm-project
 #include "mlir/Transforms/Passes.h"  // from @llvm-project
 #include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo
 #include "stablehlo/dialect/VhloOps.h"  // from @stablehlo
@@ -42,6 +47,8 @@ limitations under the License.
 #include "tensorflow/compiler/mlir/lite/debug/debug.h"
 #include "tensorflow/compiler/mlir/lite/flatbuffer_export.h"
 #include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
+#include "tensorflow/compiler/mlir/lite/python/conversion_failure_reporter.h"
+#include "tensorflow/compiler/mlir/lite/python/pass_debug_instrumentation.h"
 #include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.h"
 #include "tensorflow/compiler/mlir/lite/stablehlo/transforms/stablehlo_passes.h"
 #include "tensorflow/compiler/mlir/lite/transforms/cast_bf16_ops_to_f32_pass.h"
@@ -50,21 +57,31 @@ limitations under the License.
 #include "tensorflow/compiler/mlir/lite/transforms/optimize_broadcast_like_pass_options.h"
 #include "tensorflow/compiler/mlir/lite/transforms/pass_registry_utils.h"
 #include "tensorflow/compiler/mlir/lite/transforms/passes.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/error_util.h"
 #include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
 #include "xla/mlir_hlo/mhlo/transforms/passes.h"
 #include "xla/mlir_hlo/stablehlo_ext/transforms/passes.h"
+#include "xla/tsl/platform/env.h"
 
 namespace mlir::TFL {
 
-void AddSkipToTflitePasses(mlir::OpPassManager& pass_manager) {
+void AddPipelinePasses(mlir::OpPassManager& pass_manager,
+                       const mlir::TFL::PassConfig& pass_config) {
+  // =========================================================================
+  // 1. Skip-to-TFLite & Pre-Lowering Passes
+  // =========================================================================
   pass_manager.addNestedPass<mlir::func::FuncOp>(
       mlir::odml::CreateLegalizeChloToTflPass());
+  // Inline private functions before lowering quant annotations to eliminate
+  // func.call boundaries that cause func.call result type mismatches when
+  // quantized types are introduced.
   pass_manager.addPass(mlir::createInlinerPass());
   pass_manager.addPass(mlir::TFL::CreateLowerQuantAnnotationsPass());
   pass_manager.addPass(mlir::createSymbolDCEPass());
-}
 
-void AddHloOptimizationPasses(mlir::OpPassManager& pass_manager) {
+  // =========================================================================
+  // 2. HLO Optimization & Canonicalization Passes
+  // =========================================================================
   // Drop shape assertion custom calls before VHLO legalization
   pass_manager.addPass(mlir::odml::CreateDropShapeAssertionsPass());
 
@@ -129,16 +146,16 @@ void AddHloOptimizationPasses(mlir::OpPassManager& pass_manager) {
       mlir::createCanonicalizerPass());
   pass_manager.addNestedPass<mlir::func::FuncOp>(mlir::createCSEPass());
 
-  // Undo the MHLO::BroadcastInDimOp folding pattern on splat constants.
-  pass_manager.addPass(mlir::odml::CreateUnfoldSplatConstantPass());
-}
-
-void AddHloToTfLiteLegalizationPasses(mlir::OpPassManager& pass_manager) {
+  // =========================================================================
+  // 3. HLO to TFLite Legalization Passes
+  // =========================================================================
   // HLO -> TFLite legalization
   pass_manager.addNestedPass<mlir::func::FuncOp>(
       mlir::odml::CreateUniformQuantizedStableHloToTflPass());
   pass_manager.addNestedPass<mlir::func::FuncOp>(
       mlir::odml::CreatePrepareHloPass());
+  // This pass must be added right before the legalization because pattern
+  // rewriter driver applies folding by default.
   pass_manager.addPass(mlir::odml::CreateUnfoldSplatConstantPass());
   pass_manager.addPass(mlir::odml::CreateLegalizeHloToTfLitePass());
 
@@ -146,10 +163,10 @@ void AddHloToTfLiteLegalizationPasses(mlir::OpPassManager& pass_manager) {
   pass_manager.addPass(mlir::mhlo::createHloLegalizeToStablehloPass());
   pass_manager.addNestedPass<mlir::func::FuncOp>(
       mlir::odml::createLegalizeCompositeToCustomOpPass());
-}
 
-void AddTfLiteOptimizationPasses(mlir::OpPassManager& pass_manager,
-                                 const mlir::TFL::PassConfig& pass_config) {
+  // =========================================================================
+  // 4. TFLite Optimization & Quantization Passes
+  // =========================================================================
   pass_manager.addNestedPass<mlir::func::FuncOp>(
       mlir::TFL::CreateCastBf16OpsToF32Pass());
 
@@ -201,6 +218,69 @@ void AddTfLiteOptimizationPasses(mlir::OpPassManager& pass_manager,
   pass_manager.addPass(mlir::createReconcileUnrealizedCastsPass());
 }
 
+static absl::Status VerifyInputModule(mlir::ModuleOp module,
+                                      absl::string_view debug_dir,
+                                      bool enable_debug,
+                                      int64_t elide_elements_larger_than = 8) {
+  mlir::MLIRContext* context = module->getContext();
+  mlir::StatusScopedDiagnosticHandler initial_status_handler(
+      context,
+      /*propagate=*/false);
+  bool verification_failed = mlir::failed(module.verify());
+  absl::Status initial_status = initial_status_handler.ConsumeStatus();
+  if (verification_failed) {
+    std::string err_msg =
+        initial_status.ok()
+            ? "Input MLIR module verification failed."
+            : absl::StrCat("Input MLIR module verification failed: ",
+                           initial_status.message());
+    std::string status_code =
+        initial_status.ok() ? "INVALID_ARGUMENT"
+                            : absl::StatusCodeToString(initial_status.code());
+    ConversionFailureReporter::WriteFailureJson(
+        debug_dir, module, err_msg, "JAX_Export_Module_Verification",
+        status_code, /*failing_pass=*/"", /*failing_pass_arg=*/"",
+        /*write_module_artifacts=*/enable_debug, /*failing_function=*/"",
+        elide_elements_larger_than);
+    return absl::InvalidArgumentError(err_msg);
+  }
+  return absl::OkStatus();
+}
+
+struct PassTimingSession {
+  std::unique_ptr<std::ofstream> file_stream;
+  std::unique_ptr<llvm::raw_os_ostream> timing_stream;
+};
+
+static PassTimingSession CreatePassTimingSession(absl::string_view debug_dir) {
+  PassTimingSession session;
+  if (!tsl::Env::Default()->RecursivelyCreateDir(std::string(debug_dir)).ok()) {
+    return session;
+  }
+  std::string main_path = absl::StrCat(debug_dir, "/mlir_pass_timing.log");
+  auto file = std::make_unique<std::ofstream>(main_path);
+  if (file->is_open()) {
+    session.timing_stream = std::make_unique<llvm::raw_os_ostream>(*file);
+    session.file_stream = std::move(file);
+  }
+  return session;
+}
+
+static void AttachPassTiming(mlir::PassManager& pm,
+                             PassTimingSession& session) {
+  auto timing_manager = std::make_unique<mlir::DefaultTimingManager>();
+  timing_manager->setEnabled(true);
+  timing_manager->setDisplayMode(mlir::DefaultTimingManager::DisplayMode::List);
+
+  if (session.timing_stream) {
+    timing_manager->setOutput(mlir::createOutputStrategy(
+        mlir::DefaultTimingManager::OutputFormat::Text,
+        *session.timing_stream));
+  }
+
+  pm.enableTiming(std::move(timing_manager));
+}
+
 absl::Status ConvertStableHloToTFLite(
     mlir::ModuleOp module, const tflite::ConverterFlags& converter_flags,
     const mlir::TFL::PassConfig& pass_config,
@@ -215,29 +295,74 @@ absl::Status ConvertStableHloToTFLite(
   context->appendDialectRegistry(registry);
   context->loadAllAvailableDialects();
 
+  bool enable_debug = converter_flags.enable_debug();
+  std::string debug_dir = ConversionFailureReporter::GetOrCreateDebugDir(
+      converter_flags.debug_dir());
+
+  int64_t elide_elements_larger_than =
+      converter_flags.debug_options().has_elide_elementsattrs_if_larger()
+          ? converter_flags.debug_options().elide_elementsattrs_if_larger()
+          : 8;
+
+  if (auto status = VerifyInputModule(module, debug_dir, enable_debug,
+                                      elide_elements_larger_than);
+      !status.ok()) {
+    return status;
+  }
+
   mlir::PassManager pm(context);
-  tensorflow::InitPassManager(pm, converter_flags.debug_options(),
-                              llvm::errs());
+  PassTimingSession timing_session;
+  if (enable_debug) {
+    timing_session = CreatePassTimingSession(debug_dir);
+    AttachPassTiming(pm, timing_session);
+  }
 
-  AddSkipToTflitePasses(pm);
-  AddHloOptimizationPasses(pm);
-  AddHloToTfLiteLegalizationPasses(pm);
-  AddTfLiteOptimizationPasses(pm, pass_config);
+  tensorflow::converter::DebugOptions debug_options =
+      converter_flags.debug_options();
+  if (enable_debug) {
+    pm.getContext()->disableMultithreading();
+    debug_options.clear_print_ir_before();
+    debug_options.clear_print_ir_after();
+  }
+  tensorflow::InitPassManager(pm, debug_options, llvm::nulls());
+
+  AddPipelinePasses(pm, pass_config);
+
+  PipelineFailureCoordinator failure_coordinator(debug_dir, enable_debug,
+                                                 elide_elements_larger_than);
+  pm.addInstrumentation(failure_coordinator.CreateInstrumentation(
+      converter_flags.debug_options().print_ir_before(),
+      converter_flags.debug_options().print_ir_after()));
+
+  mlir::StatusScopedDiagnosticHandler status_handler(context,
+                                                     /*propagate=*/false);
+  bool pass_failed = mlir::failed(pm.run(module));
+  absl::Status pass_status = status_handler.ConsumeStatus();
 
-  if (mlir::failed(pm.run(module))) {
-    return absl::InvalidArgumentError("StableHLO to TFLite pipeline failed.");
+  if (timing_session.timing_stream) {
+    timing_session.timing_stream->flush();
+  }
+  if (timing_session.file_stream) {
+    timing_session.file_stream->flush();
+  }
+
+  if (pass_failed) {
+    return failure_coordinator.ReportFailure(module, pass_status);
   }
 
   tflite::FlatbufferExportOptions options;
   options.converter_flags.set_allow_custom_ops(true);
   options.converter_flags.set_use_buffer_offset(true);
 
+  std::string diag_errors;
+  SerializationDiagHandler diag_handler(module.getContext(), &diag_errors);
+
   auto status =
       tflite::MlirToFlatBufferTranslateFunction(module, options, export_stream);
 
   if (!status.ok()) {
-    return absl::InvalidArgumentError(
-        absl::StrCat("Failed to serialize to FlatBuffer: ", status.message()));
+    return failure_coordinator.ReportSerializationFailure(module, status,
+                                                          diag_errors);
   }
 
   return absl::OkStatus();
diff --git a/tensorflow/core/framework/metrics.cc b/tensorflow/core/framework/metrics.cc
index 097f5c286c9f41..318c71af0e1d13 100644
--- a/tensorflow/core/framework/metrics.cc
+++ b/tensorflow/core/framework/metrics.cc
@@ -28,6 +28,7 @@ limitations under the License.
 #include "xla/tsl/lib/monitoring/sampler.h"
 #include "xla/tsl/platform/types.h"
 #include "xla/tsl/protobuf/error_codes.pb.h"
+#include "tensorflow/core/platform/env.h"
 #include "tensorflow/core/protobuf/data_service.pb.h"
 
 namespace tensorflow {
@@ -445,6 +446,14 @@ auto* xla_compilation_time_usecs = tsl::monitoring::Counter<0>::New(
     "/tensorflow/core/xla_compilation_time_usecs",
     "The total time spent on compiling XLA graphs in microseconds.");
 
+auto* xla_compilation_start_time = tsl::monitoring::Gauge<int64_t, 0>::New(
+    "/tensorflow/core/xla_compilation_start_time",
+    "Timestamp of when the most recent XLA compilation started.");
+
+auto* xla_compilation_end_time = tsl::monitoring::Gauge<int64_t, 0>::New(
+    "/tensorflow/core/xla_compilation_end_time",
+    "Timestamp of when the most recent XLA compilation ended.");
+
 auto* xla_tpu_spmd_cores_per_replica = tsl::monitoring::Counter<1>::New(
     "/tensorflow/tpu/xla_spmd_cores_per_replica",
     "The number of cores used by XLA SPMD-replicated models.", "cores");
@@ -972,7 +981,12 @@ void UpdateTpuVariableDistributionTime(const uint64_t distribution_time_usecs) {
   }
 }
 
-void UpdateXlaCompilationTime(const uint64_t compilation_time_usecs) {
+void UpdateXlaCompilationStartTime(const uint64_t compilation_start_time_us) {
+  xla_compilation_start_time->GetCell()->Set(compilation_start_time_us);
+}
+
+void UpdateXlaCompilationTime(const uint64_t compilation_time_usecs,
+                              const uint64_t compile_end_us) {
   if (compilation_time_usecs > 0) {
     static auto* xla_compilations_cell = xla_compilations->GetCell();
     static auto* xla_compilation_time_usecs_cell =
@@ -980,6 +994,9 @@ void UpdateXlaCompilationTime(const uint64_t compilation_time_usecs) {
     xla_compilations_cell->IncrementBy(1);
     xla_compilation_time_usecs_cell->IncrementBy(compilation_time_usecs);
   }
+  uint64_t final_end_us =
+      (compile_end_us > 0) ? compile_end_us : Env::Default()->NowMicros();
+  xla_compilation_end_time->GetCell()->Set(final_end_us);
 }
 
 void RecordUnusedOutput(const std::string& op_name) {
diff --git a/tensorflow/core/framework/metrics.h b/tensorflow/core/framework/metrics.h
index 7fcbeae1b811b6..846169c4e8e242 100644
--- a/tensorflow/core/framework/metrics.h
+++ b/tensorflow/core/framework/metrics.h
@@ -543,7 +543,9 @@ monitoring::Counter<2>* GetGraphOptimizationCounter();
 void UpdateTpuVariableDistributionTime(const uint64_t distribution_time_usecs);
 
 // Updates the metrics stored about time XLA spents compiling graphs.
-void UpdateXlaCompilationTime(const uint64_t compilation_time_usecs);
+void UpdateXlaCompilationTime(uint64_t compilation_time_usecs,
+                              uint64_t compile_end_us = 0);
+void UpdateXlaCompilationStartTime(uint64_t compilation_start_time_us);
 
 // Increments (by 1) a simple integer counter that is exposed for testing.
 void IncrementTestCounter(const std::string& name, const std::string& label);
diff --git a/tensorflow/core/framework/metrics_test.cc b/tensorflow/core/framework/metrics_test.cc
index fe8ea8056e4c37..cb99488d8e17a8 100644
--- a/tensorflow/core/framework/metrics_test.cc
+++ b/tensorflow/core/framework/metrics_test.cc
@@ -158,4 +158,21 @@ TEST(Metrics, TFDataPrefetchBufferSize) {
   EXPECT_EQ(gauge.Read("node_1"), 3);
 }
 
+TEST(Metrics, UpdateXlaCompilationTime) {
+  CellReader<int64_t> start_counter(
+      "/tensorflow/core/xla_compilation_start_time");
+  CellReader<int64_t> end_counter("/tensorflow/core/xla_compilation_end_time");
+  CellReader<int64_t> counter("/tensorflow/core/xla_compilations");
+  CellReader<int64_t> time_counter(
+      "/tensorflow/core/xla_compilation_time_usecs");
+
+  tensorflow::metrics::UpdateXlaCompilationStartTime(100);
+  tensorflow::metrics::UpdateXlaCompilationTime(500, 600);
+
+  EXPECT_EQ(start_counter.Read(), 100);
+  EXPECT_EQ(end_counter.Read(), 600);
+  EXPECT_EQ(counter.Read(), 1);
+  EXPECT_EQ(time_counter.Read(), 500);
+}
+
 }  // namespace
diff --git a/third_party/xla/third_party/tsl/tsl/profiler/lib/continuous_profiler_orchestrator.h b/third_party/xla/third_party/tsl/tsl/profiler/lib/continuous_profiler_orchestrator.h
index d12a8be24acb95..58954be377cb31 100644
--- a/third_party/xla/third_party/tsl/tsl/profiler/lib/continuous_profiler_orchestrator.h
+++ b/third_party/xla/third_party/tsl/tsl/profiler/lib/continuous_profiler_orchestrator.h
@@ -18,6 +18,7 @@ limitations under the License.
 #include <algorithm>
 #include <any>
 #include <cstddef>
+#include <cstdint>
 #include <deque>
 #include <memory>
 #include <utility>
@@ -37,6 +38,15 @@ limitations under the License.
 namespace tsl {
 namespace profiler {
 
+inline constexpr size_t kDefaultMaxBufferBytes =
+    4ULL * 1024 * 1024 * 1024;  // 4GB
+
+struct DrainedBuffer {
+  std::vector<std::any> chunks;
+  uint64_t cumulative_dropped_chunks = 0;
+  uint64_t cumulative_dropped_bytes = 0;
+};
+
 template <typename ProfilerType>
 class ContinuousProfilerOrchestrator : public ProfilerInterface {
  public:
@@ -45,8 +55,10 @@ class ContinuousProfilerOrchestrator : public ProfilerInterface {
   static constexpr absl::Duration kMaxPollingInterval = absl::Seconds(5);
 
   explicit ContinuousProfilerOrchestrator(
-      std::unique_ptr<ProfilerType> profiler)
+      std::unique_ptr<ProfilerType> profiler,
+      size_t max_buffer_bytes = kDefaultMaxBufferBytes)
       : profiler_(std::move(profiler)),
+        max_buffer_bytes_(max_buffer_bytes),
         is_running_(false),
         polling_interval_(kDefaultPollingInterval) {}
 
@@ -77,10 +89,10 @@ class ContinuousProfilerOrchestrator : public ProfilerInterface {
   // Stops background thread and profiling.
   absl::Status Stop() override {
     absl::Status status = StopInternal();
-    auto result = profiler_->Consume();
+    absl::StatusOr<ConsumeResult> result = profiler_->Consume();
     if (result.ok()) {
       absl::MutexLock lock(mutex_);
-      circular_buffer_.push_back(std::move(result->data));
+      PushChunkLocked(std::move(*result));
     } else if (!absl::IsUnimplemented(result.status())) {
       LOG(WARNING) << "Final Consume failed during Stop: " << result.status();
     }
@@ -129,18 +141,87 @@ class ContinuousProfilerOrchestrator : public ProfilerInterface {
   ProfilerType* profiler() { return profiler_.get(); }
   const ProfilerType* profiler() const { return profiler_.get(); }
 
-  std::vector<std::any> PopBuffer() {
+  size_t max_buffer_bytes() const { return max_buffer_bytes_; }
+
+  size_t total_buffered_bytes() const {
     absl::MutexLock lock(mutex_);
+    return total_buffered_bytes_;
+  }
+
+  uint64_t dropped_chunks_count() const {
+    absl::MutexLock lock(mutex_);
+    return dropped_chunks_count_;
+  }
+
+  uint64_t dropped_bytes_count() const {
+    absl::MutexLock lock(mutex_);
+    return dropped_bytes_count_;
+  }
+
+  DrainedBuffer PopBufferWithTelemetry() {
+    std::deque<std::any> local_buffer;
+    uint64_t dropped_chunks = 0;
+    uint64_t dropped_bytes = 0;
+    {
+      absl::MutexLock lock(mutex_);
+      local_buffer.swap(circular_buffer_);
+      chunk_sizes_.clear();
+      total_buffered_bytes_ = 0;
+      dropped_chunks = dropped_chunks_count_;
+      dropped_bytes = dropped_bytes_count_;
+    }
+
     std::vector<std::any> chunks;
-    chunks.reserve(circular_buffer_.size());
-    for (auto& item : circular_buffer_) {
-      chunks.push_back(std::move(item));
+    chunks.reserve(local_buffer.size());
+    for (auto& item : local_buffer) {
+      if (item.has_value()) {
+        chunks.push_back(std::move(item));
+      }
     }
-    circular_buffer_.clear();
-    return chunks;
+    return DrainedBuffer{
+        .chunks = std::move(chunks),
+        .cumulative_dropped_chunks = dropped_chunks,
+        .cumulative_dropped_bytes = dropped_bytes,
+    };
   }
 
+  std::vector<std::any> PopBuffer() { return PopBufferWithTelemetry().chunks; }
+
  private:
+  void PushChunkLocked(ConsumeResult chunk)
+      ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_) {
+    if (!chunk.data.has_value()) {
+      return;
+    }
+
+    if (chunk.estimated_size_bytes > max_buffer_bytes_) {
+      LOG_EVERY_N_SEC(WARNING, 30)
+          << "ContinuousProfilerOrchestrator rejected chunk of size "
+          << chunk.estimated_size_bytes << " bytes exceeding max buffer limit "
+          << max_buffer_bytes_ << " bytes.";
+      dropped_chunks_count_ += 1;
+      dropped_bytes_count_ += chunk.estimated_size_bytes;
+      return;
+    }
+
+    while (!circular_buffer_.empty() &&
+           (total_buffered_bytes_ + chunk.estimated_size_bytes >
+            max_buffer_bytes_)) {
+      size_t front_size = chunk_sizes_.front();
+      total_buffered_bytes_ = (total_buffered_bytes_ > front_size)
+                                  ? total_buffered_bytes_ - front_size
+                                  : 0;
+      dropped_chunks_count_ += 1;
+      dropped_bytes_count_ += front_size;
+      circular_buffer_.pop_front();
+      chunk_sizes_.pop_front();
+    }
+
+    total_buffered_bytes_ += chunk.estimated_size_bytes;
+    circular_buffer_.push_back(std::move(chunk.data));
+    chunk_sizes_.push_back(chunk.estimated_size_bytes);
+  }
+
   void IngestionLoop() {
     LOG(INFO) << "ContinuousProfilerOrchestrator::IngestionLoop started";
     while (true) {
@@ -152,14 +233,9 @@ class ContinuousProfilerOrchestrator : public ProfilerInterface {
 
       absl::MutexLock lock(mutex_);
       if (result.ok()) {
-        circular_buffer_.push_back(std::move(result->data));
-
-        // Cap circular buffer to prevent infinite memory growth.
-        if (circular_buffer_.size() > 100) {
-          circular_buffer_.pop_front();
-        }
-
-        AdjustIntervalLocked(result->estimated_size_bytes);
+        const size_t chunk_size = result->estimated_size_bytes;
+        PushChunkLocked(std::move(*result));
+        AdjustIntervalLocked(chunk_size);
       }
 
       if (!is_running_) break;
@@ -195,6 +271,7 @@ class ContinuousProfilerOrchestrator : public ProfilerInterface {
   }
 
   std::unique_ptr<ProfilerType> profiler_;
+  const size_t max_buffer_bytes_;
 
   mutable absl::Mutex mutex_;
   absl::CondVar cv_;
@@ -203,6 +280,10 @@ class ContinuousProfilerOrchestrator : public ProfilerInterface {
 
   absl::Duration polling_interval_ ABSL_GUARDED_BY(mutex_);
   std::deque<std::any> circular_buffer_ ABSL_GUARDED_BY(mutex_);
+  std::deque<size_t> chunk_sizes_ ABSL_GUARDED_BY(mutex_);
+  size_t total_buffered_bytes_ ABSL_GUARDED_BY(mutex_) = 0;
+  uint64_t dropped_chunks_count_ ABSL_GUARDED_BY(mutex_) = 0;
+  uint64_t dropped_bytes_count_ ABSL_GUARDED_BY(mutex_) = 0;
 };
 
 }  // namespace profiler
diff --git a/third_party/xla/third_party/tsl/tsl/profiler/lib/continuous_profiler_orchestrator_test.cc b/third_party/xla/third_party/tsl/tsl/profiler/lib/continuous_profiler_orchestrator_test.cc
index a1c81fde6b7cc3..4341e5eeb95f18 100644
--- a/third_party/xla/third_party/tsl/tsl/profiler/lib/continuous_profiler_orchestrator_test.cc
+++ b/third_party/xla/third_party/tsl/tsl/profiler/lib/continuous_profiler_orchestrator_test.cc
@@ -77,7 +77,7 @@ TEST(ContinuousProfilerOrchestratorTest,
         }
         return ConsumeResult{
             .data = std::any(count),
-            .estimated_size_bytes = 1000 * 1024 * 1024  // 1000MB (>512MB)
+            .estimated_size_bytes = 600 * 1024 * 1024  // 600MB (>512MB)
         };
       });
 
@@ -186,6 +186,314 @@ TEST(ContinuousProfilerOrchestratorTest, SerializeChunks) {
   EXPECT_EQ(spaces.size(), 1);
 }
 
+TEST(ContinuousProfilerOrchestratorTest, CustomMemoryBudget) {
+  auto mock_profiler = std::make_unique<MockProfiler>();
+  ContinuousProfilerOrchestrator<ProfilerInterface> default_orchestrator(
+      std::move(mock_profiler));
+  EXPECT_EQ(default_orchestrator.max_buffer_bytes(), kDefaultMaxBufferBytes);
+
+  auto mock_profiler2 = std::make_unique<MockProfiler>();
+  ContinuousProfilerOrchestrator<ProfilerInterface> custom_orchestrator(
+      std::move(mock_profiler2), 500 * 1024 * 1024);
+  EXPECT_EQ(custom_orchestrator.max_buffer_bytes(), 500 * 1024 * 1024);
+}
+
+TEST(ContinuousProfilerOrchestratorTest, ByteBudgetAccounting) {
+  auto mock_profiler = std::make_unique<MockProfiler>();
+  MockProfiler* mock = mock_profiler.get();
+
+  EXPECT_CALL(*mock, Start()).WillOnce(Return(absl::OkStatus()));
+  EXPECT_CALL(*mock, Stop()).WillOnce(Return(absl::OkStatus()));
+
+  absl::Notification consumed_2;
+  std::atomic<int> consume_count(0);
+  EXPECT_CALL(*mock, Consume())
+      .WillRepeatedly([&]() -> absl::StatusOr<ConsumeResult> {
+        int count = ++consume_count;
+        if (count <= 2) {
+          if (count == 2 && !consumed_2.HasBeenNotified()) {
+            consumed_2.Notify();
+          }
+          return ConsumeResult{
+              .data = std::any(count),
+              .estimated_size_bytes = 25 * 1024 * 1024,  // 25MB
+          };
+        }
+        return absl::OutOfRangeError("End of stream");
+      });
+
+  ContinuousProfilerOrchestrator<ProfilerInterface> orchestrator(
+      std::move(mock_profiler), 100 * 1024 * 1024);
+  EXPECT_EQ(orchestrator.total_buffered_bytes(), 0);
+  EXPECT_EQ(orchestrator.dropped_chunks_count(), 0);
+  EXPECT_EQ(orchestrator.dropped_bytes_count(), 0);
+
+  ASSERT_OK(orchestrator.Start());
+  consumed_2.WaitForNotification();
+  ASSERT_OK(orchestrator.Stop());
+
+  EXPECT_EQ(orchestrator.total_buffered_bytes(), 50 * 1024 * 1024);
+  EXPECT_EQ(orchestrator.dropped_chunks_count(), 0);
+  EXPECT_EQ(orchestrator.dropped_bytes_count(), 0);
+
+  std::vector<std::any> chunks = orchestrator.PopBuffer();
+  EXPECT_EQ(chunks.size(), 2);
+  EXPECT_EQ(orchestrator.total_buffered_bytes(), 0);
+}
+
+TEST(ContinuousProfilerOrchestratorTest, FIFOEvictionOnMemoryCap) {
+  auto mock_profiler = std::make_unique<MockProfiler>();
+  MockProfiler* mock = mock_profiler.get();
+
+  EXPECT_CALL(*mock, Start()).WillOnce(Return(absl::OkStatus()));
+  EXPECT_CALL(*mock, Stop()).WillOnce(Return(absl::OkStatus()));
+
+  absl::Notification consumed_3;
+  std::atomic<int> consume_count(0);
+  EXPECT_CALL(*mock, Consume())
+      .WillRepeatedly([&]() -> absl::StatusOr<ConsumeResult> {
+        int count = ++consume_count;
+        if (count <= 3) {
+          if (count == 3 && !consumed_3.HasBeenNotified()) {
+            consumed_3.Notify();
+          }
+          return ConsumeResult{
+              .data = std::any(count),
+              .estimated_size_bytes = 40 * 1024 * 1024,  // 40MB
+          };
+        }
+        return absl::OutOfRangeError("End of stream");
+      });
+
+  ContinuousProfilerOrchestrator<ProfilerInterface> orchestrator(
+      std::move(mock_profiler), 100 * 1024 * 1024);  // 100MB limit
+
+  ASSERT_OK(orchestrator.Start());
+  consumed_3.WaitForNotification();
+  ASSERT_OK(orchestrator.Stop());
+
+  // Chunks: 1 (40MB), 2 (40MB), 3 (40MB). Total = 120MB > 100MB limit.
+  // Chunk 1 is evicted. Retained: chunks 2 & 3 (80MB).
+  EXPECT_EQ(orchestrator.total_buffered_bytes(), 80 * 1024 * 1024);
+  EXPECT_EQ(orchestrator.dropped_chunks_count(), 1);
+  EXPECT_EQ(orchestrator.dropped_bytes_count(), 40 * 1024 * 1024);
+
+  DrainedBuffer drained = orchestrator.PopBufferWithTelemetry();
+  ASSERT_EQ(drained.chunks.size(), 2);
+  EXPECT_EQ(std::any_cast<int>(drained.chunks[0]), 2);
+  EXPECT_EQ(std::any_cast<int>(drained.chunks[1]), 3);
+  EXPECT_EQ(drained.cumulative_dropped_chunks, 1);
+  EXPECT_EQ(drained.cumulative_dropped_bytes, 40 * 1024 * 1024);
+
+  // Buffer is empty after drain, but cumulative drop counters persist
+  EXPECT_EQ(orchestrator.total_buffered_bytes(), 0);
+  EXPECT_EQ(orchestrator.dropped_chunks_count(), 1);
+  EXPECT_EQ(orchestrator.dropped_bytes_count(), 40 * 1024 * 1024);
+}
+
+TEST(ContinuousProfilerOrchestratorTest, OversizedChunkRejection) {
+  auto mock_profiler = std::make_unique<MockProfiler>();
+  MockProfiler* mock = mock_profiler.get();
+
+  EXPECT_CALL(*mock, Start()).WillOnce(Return(absl::OkStatus()));
+  EXPECT_CALL(*mock, Stop()).WillOnce(Return(absl::OkStatus()));
+
+  absl::Notification consumed_1;
+  std::atomic<int> consume_count(0);
+  EXPECT_CALL(*mock, Consume())
+      .WillRepeatedly([&]() -> absl::StatusOr<ConsumeResult> {
+        int count = ++consume_count;
+        if (count == 1) {
+          consumed_1.Notify();
+          return ConsumeResult{
+              .data = std::any(count),
+              .estimated_size_bytes = 60 * 1024 * 1024,  // 60MB > 50MB limit
+          };
+        }
+        return absl::OutOfRangeError("End of stream");
+      });
+
+  ContinuousProfilerOrchestrator<ProfilerInterface> orchestrator(
+      std::move(mock_profiler), 50 * 1024 * 1024);
+
+  ASSERT_OK(orchestrator.Start());
+  consumed_1.WaitForNotification();
+  ASSERT_OK(orchestrator.Stop());
+
+  EXPECT_EQ(orchestrator.total_buffered_bytes(), 0);
+  EXPECT_EQ(orchestrator.dropped_chunks_count(), 1);
+  EXPECT_EQ(orchestrator.dropped_bytes_count(), 60 * 1024 * 1024);
+  EXPECT_TRUE(orchestrator.PopBuffer().empty());
+}
+
+TEST(ContinuousProfilerOrchestratorTest, IdleEmptyDataChunkNoOp) {
+  auto mock_profiler = std::make_unique<MockProfiler>();
+  MockProfiler* mock = mock_profiler.get();
+
+  EXPECT_CALL(*mock, Start()).WillOnce(Return(absl::OkStatus()));
+  EXPECT_CALL(*mock, Stop()).WillOnce(Return(absl::OkStatus()));
+
+  absl::Notification consumed_1;
+  std::atomic<int> consume_count(0);
+  EXPECT_CALL(*mock, Consume())
+      .WillRepeatedly([&]() -> absl::StatusOr<ConsumeResult> {
+        int count = ++consume_count;
+        if (count == 1) {
+          consumed_1.Notify();
+          return ConsumeResult{
+              .data = std::any(),  // empty data (no value)
+              .estimated_size_bytes = 0,
+          };
+        }
+        return absl::OutOfRangeError("End of stream");
+      });
+
+  ContinuousProfilerOrchestrator<ProfilerInterface> orchestrator(
+      std::move(mock_profiler));
+
+  ASSERT_OK(orchestrator.Start());
+  consumed_1.WaitForNotification();
+  ASSERT_OK(orchestrator.Stop());
+
+  EXPECT_EQ(orchestrator.total_buffered_bytes(), 0);
+  EXPECT_EQ(orchestrator.dropped_chunks_count(), 0);
+  EXPECT_EQ(orchestrator.dropped_bytes_count(), 0);
+  EXPECT_TRUE(orchestrator.PopBuffer().empty());
+}
+
+TEST(ContinuousProfilerOrchestratorTest, ZeroByteChunkWithDataIsBuffered) {
+  auto mock_profiler = std::make_unique<MockProfiler>();
+  MockProfiler* mock = mock_profiler.get();
+
+  EXPECT_CALL(*mock, Start()).WillOnce(Return(absl::OkStatus()));
+  EXPECT_CALL(*mock, Stop()).WillOnce(Return(absl::OkStatus()));
+
+  absl::Notification consumed_1;
+  std::atomic<int> consume_count(0);
+  EXPECT_CALL(*mock, Consume())
+      .WillRepeatedly([&]() -> absl::StatusOr<ConsumeResult> {
+        int count = ++consume_count;
+        if (count == 1) {
+          consumed_1.Notify();
+          return ConsumeResult{
+              .data = std::any(42),       // valid payload
+              .estimated_size_bytes = 0,  // 0 estimated bytes
+          };
+        }
+        return absl::OutOfRangeError("End of stream");
+      });
+
+  ContinuousProfilerOrchestrator<ProfilerInterface> orchestrator(
+      std::move(mock_profiler));
+
+  ASSERT_OK(orchestrator.Start());
+  consumed_1.WaitForNotification();
+  ASSERT_OK(orchestrator.Stop());
+
+  EXPECT_EQ(orchestrator.total_buffered_bytes(), 0);
+  EXPECT_EQ(orchestrator.dropped_chunks_count(), 0);
+  EXPECT_EQ(orchestrator.dropped_bytes_count(), 0);
+  std::vector<std::any> chunks = orchestrator.PopBuffer();
+  ASSERT_EQ(chunks.size(), 1);
+  EXPECT_EQ(std::any_cast<int>(chunks[0]), 42);
+}
+
+TEST(ContinuousProfilerOrchestratorTest, DropCountersAccuracy) {
+  auto mock_profiler = std::make_unique<MockProfiler>();
+  MockProfiler* mock = mock_profiler.get();
+
+  EXPECT_CALL(*mock, Start()).WillOnce(Return(absl::OkStatus()));
+  EXPECT_CALL(*mock, Stop()).WillOnce(Return(absl::OkStatus()));
+
+  absl::Notification consumed_3;
+  std::atomic<int> consume_count(0);
+  EXPECT_CALL(*mock, Consume())
+      .WillRepeatedly([&]() -> absl::StatusOr<ConsumeResult> {
+        int count = ++consume_count;
+        if (count == 1) {
+          return ConsumeResult{
+              .data = std::any(1),
+              .estimated_size_bytes = 60 * 1024 * 1024,  // 60MB
+          };
+        }
+        if (count == 2) {
+          return ConsumeResult{
+              .data = std::any(2),
+              .estimated_size_bytes = 120 * 1024 * 1024,  // 120MB (rejected)
+          };
+        }
+        if (count == 3) {
+          consumed_3.Notify();
+          return ConsumeResult{
+              .data = std::any(3),
+              .estimated_size_bytes =
+                  60 * 1024 * 1024,  // 60MB (evicts chunk 1)
+          };
+        }
+        return absl::OutOfRangeError("End of stream");
+      });
+
+  ContinuousProfilerOrchestrator<ProfilerInterface> orchestrator(
+      std::move(mock_profiler), 100 * 1024 * 1024);
+
+  ASSERT_OK(orchestrator.Start());
+  consumed_3.WaitForNotification();
+  ASSERT_OK(orchestrator.Stop());
+
+  EXPECT_EQ(orchestrator.total_buffered_bytes(), 60 * 1024 * 1024);
+  EXPECT_EQ(orchestrator.dropped_chunks_count(), 2);
+  EXPECT_EQ(orchestrator.dropped_bytes_count(), (120 + 60) * 1024 * 1024);
+
+  DrainedBuffer drained = orchestrator.PopBufferWithTelemetry();
+  ASSERT_EQ(drained.chunks.size(), 1);
+  EXPECT_EQ(std::any_cast<int>(drained.chunks[0]), 3);
+  EXPECT_EQ(drained.cumulative_dropped_chunks, 2);
+  EXPECT_EQ(drained.cumulative_dropped_bytes, 180 * 1024 * 1024);
+}
+
+TEST(ContinuousProfilerOrchestratorTest, StopDrainAccounting) {
+  auto mock_profiler = std::make_unique<MockProfiler>();
+  MockProfiler* mock = mock_profiler.get();
+
+  EXPECT_CALL(*mock, Start()).WillOnce(Return(absl::OkStatus()));
+  EXPECT_CALL(*mock, Stop()).WillOnce(Return(absl::OkStatus()));
+
+  absl::Notification consumed_1;
+  std::atomic<int> consume_count(0);
+  EXPECT_CALL(*mock, Consume())
+      .WillRepeatedly([&]() -> absl::StatusOr<ConsumeResult> {
+        int count = ++consume_count;
+        if (count == 1) {
+          consumed_1.Notify();
+          return ConsumeResult{
+              .data = std::any(1),
+              .estimated_size_bytes = 20 * 1024 * 1024,  // 20MB
+          };
+        }
+        if (count == 2) {
+          // Returned during Stop()
+          return ConsumeResult{
+              .data = std::any(2),
+              .estimated_size_bytes = 30 * 1024 * 1024,  // 30MB
+          };
+        }
+        return absl::OutOfRangeError("End of stream");
+      });
+
+  ContinuousProfilerOrchestrator<ProfilerInterface> orchestrator(
+      std::move(mock_profiler), 100 * 1024 * 1024);
+
+  ASSERT_OK(orchestrator.Start());
+  consumed_1.WaitForNotification();
+  ASSERT_OK(orchestrator.Stop());
+
+  EXPECT_EQ(orchestrator.total_buffered_bytes(), 50 * 1024 * 1024);
+  std::vector<std::any> chunks = orchestrator.PopBuffer();
+  ASSERT_EQ(chunks.size(), 2);
+  EXPECT_EQ(std::any_cast<int>(chunks[0]), 1);
+  EXPECT_EQ(std::any_cast<int>(chunks[1]), 2);
+}
+
 }  // namespace
 }  // namespace profiler
 }  // namespace tsl
diff --git a/third_party/xla/xla/backends/autotuner/BUILD b/third_party/xla/xla/backends/autotuner/BUILD
index 9d0c34d30a8140..23945aaf4543f1 100644
--- a/third_party/xla/xla/backends/autotuner/BUILD
+++ b/third_party/xla/xla/backends/autotuner/BUILD
@@ -180,6 +180,7 @@ cc_library(
         ":autotune_fingerprint",
         ":autotuner_cache_interface",
         ":autotuning_proto_cc",
+        ":backends_proto_cc",
         ":codegen_orchestrator",
         ":config_runner",
         ":config_selector",
@@ -554,12 +555,15 @@ cc_library(
     srcs = ["config_selector.cc"],
     hdrs = ["config_selector.h"],
     deps = [
+        ":backends_proto_cc",
         ":config_runner",
+        "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/time",
+        "@com_google_absl//absl/types:span",
     ],
 )
 
diff --git a/third_party/xla/xla/backends/autotuner/autotuner.cc b/third_party/xla/xla/backends/autotuner/autotuner.cc
index 15ace17c0b8454..c30538f659345d 100644
--- a/third_party/xla/xla/backends/autotuner/autotuner.cc
+++ b/third_party/xla/xla/backends/autotuner/autotuner.cc
@@ -238,7 +238,8 @@ tsl::Future<Autotuner::Config> Autotuner::GetTunedConfig(
 
         ABSL_ASSIGN_OR_RETURN(
             ConfigRunner::ConfigProfile best_profile,
-            PickBestConfig(profiles, options_.scratch_bytes_window_size_us));
+            PickBestConfig(profiles, options_.scratch_bytes_window_size_us,
+                           options_.excluded_backends));
 
         return std::move(best_profile.config);
       });
diff --git a/third_party/xla/xla/backends/autotuner/autotuner.h b/third_party/xla/xla/backends/autotuner/autotuner.h
index da850d0ab54790..adf7d29d5d1b56 100644
--- a/third_party/xla/xla/backends/autotuner/autotuner.h
+++ b/third_party/xla/xla/backends/autotuner/autotuner.h
@@ -30,6 +30,7 @@ limitations under the License.
 #include "xla/autotune_results.pb.h"
 #include "xla/backends/autotuner/autotuner_cache_interface.h"
 #include "xla/backends/autotuner/autotuning.pb.h"
+#include "xla/backends/autotuner/backends.pb.h"
 #include "xla/backends/autotuner/codegen_orchestrator.h"
 #include "xla/backends/autotuner/config_runner.h"
 #include "xla/backends/autotuner/hlo_extractor.h"
@@ -47,6 +48,7 @@ class Autotuner {
  public:
   struct Options {
     int scratch_bytes_window_size_us = 2;
+    std::vector<autotuner::Backend> excluded_backends;
     ConfigRunner::CorrectnessCheckOptions correctness_check_options;
     // File path to dump the profiles for all configs profiled for each HLO
     // instruction.
diff --git a/third_party/xla/xla/backends/autotuner/autotuner_test.cc b/third_party/xla/xla/backends/autotuner/autotuner_test.cc
index e129f9cc0034f2..ae7f3f547dbc8c 100644
--- a/third_party/xla/xla/backends/autotuner/autotuner_test.cc
+++ b/third_party/xla/xla/backends/autotuner/autotuner_test.cc
@@ -814,5 +814,95 @@ TEST_F(AutotunerTest, DumpLogsWithCacheContext) {
   EXPECT_EQ(actual_profiles.instruction_profiles_size(), 1);
 }
 
+TEST_F(AutotunerTest, ExcludedBackendsAreProfiledForReferenceButNotPicked) {
+  Autotuner::Options options;
+  options.excluded_backends = {autotuner::Backend::CUBLASLT_FISSION};
+  options.correctness_check_options.enable_correctness_check = true;
+
+  auto cublas_backend = std::make_unique<MockCodegenBackend>();
+  EXPECT_CALL(*cublas_backend, name()).WillRepeatedly(Return("cublas_backend"));
+  EXPECT_CALL(*cublas_backend, backend())
+      .WillRepeatedly(Return(autotuner::Backend::CUBLASLT_FISSION));
+  EXPECT_CALL(*cublas_backend, CanProduceWrongResults())
+      .WillRepeatedly(Return(false));
+  std::vector<std::unique_ptr<BackendConfig>> cublas_configs;
+  cublas_configs.push_back(GetTestConfig("cublas_config"));
+  EXPECT_CALL(*cublas_backend, GetSupportedConfigs)
+      .WillOnce(Return(std::move(cublas_configs)));
+  EXPECT_CALL(*cublas_backend, Compile(_, _)).WillRepeatedly([] {
+    return std::unique_ptr<Executable>();
+  });
+
+  auto triton_backend = std::make_unique<MockCodegenBackend>();
+  EXPECT_CALL(*triton_backend, name()).WillRepeatedly(Return("triton_backend"));
+  EXPECT_CALL(*triton_backend, backend())
+      .WillRepeatedly(Return(autotuner::Backend::TRITON));
+  EXPECT_CALL(*triton_backend, CanProduceWrongResults())
+      .WillRepeatedly(Return(true));
+  std::vector<std::unique_ptr<BackendConfig>> triton_configs;
+  triton_configs.push_back(GetTestConfig("triton_config"));
+  EXPECT_CALL(*triton_backend, GetSupportedConfigs)
+      .WillOnce(Return(std::move(triton_configs)));
+  EXPECT_CALL(*triton_backend, Compile(_, _)).WillRepeatedly([] {
+    return std::unique_ptr<Executable>();
+  });
+
+  auto profiler = std::make_unique<MockProfiler>();
+  EXPECT_CALL(*profiler, CreateInputBuffers(_, _))
+      .WillOnce(Return(std::make_unique<InputBuffers>()));
+  EXPECT_CALL(*profiler, CheckInputBuffers(_))
+      .WillRepeatedly(Return(absl::OkStatus()));
+  EXPECT_CALL(*profiler, CheckOutputBuffer(_, _, _))
+      .WillRepeatedly(Return(absl::OkStatus()));
+
+  // cuBLAS duration 50us (faster), Triton duration 100us.
+  // Both return valid outputs and match reference.
+  EXPECT_CALL(*profiler, Profile(_, _))
+      .WillOnce([] {
+        ProfileResult res({absl::Microseconds(50)});
+        res.output_buffer =
+            ScopedShapedBuffer(ShapeUtil::MakeShape(F32, {4}), nullptr, 0);
+        return res;
+      })
+      .WillOnce([] {
+        ProfileResult res({absl::Microseconds(100)});
+        res.output_buffer =
+            ScopedShapedBuffer(ShapeUtil::MakeShape(F32, {4}), nullptr, 0);
+        return res;
+      });
+
+  std::vector<std::unique_ptr<CodegenBackend>> backends;
+  backends.push_back(std::move(cublas_backend));
+  backends.push_back(std::move(triton_backend));
+
+  ASSERT_OK_AND_ASSIGN(auto orchestrator,
+                       CodegenOrchestrator::Create(std::move(backends), {}));
+
+  std::vector<std::unique_ptr<Profiler>> profilers;
+  profilers.push_back(std::move(profiler));
+
+  ASSERT_OK_AND_ASSIGN(auto autotuner,
+                       Autotuner::Create(std::move(orchestrator),
+                                         std::move(profilers), options));
+
+  constexpr absl::string_view kHlo = R"(
+    HloModule test_module
+    ENTRY main {
+      p0 = f32[4] parameter(0)
+      ROOT copy = f32[4] copy(p0)
+    }
+  )";
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kHlo));
+
+  ASSERT_OK_AND_ASSIGN(
+      auto results,
+      autotuner->TuneConfigs(*module, [](const HloInstruction& instr) {
+        return instr.opcode() == HloOpcode::kCopy;
+      }));
+  ASSERT_EQ(results.size(), 1);
+  EXPECT_EQ(results[0].config.codegen_backend->backend(),
+            autotuner::Backend::TRITON);
+}
+
 }  // namespace
 }  // namespace xla
diff --git a/third_party/xla/xla/backends/autotuner/codegen_orchestrator.cc b/third_party/xla/xla/backends/autotuner/codegen_orchestrator.cc
index 42036c4237bc76..27c084b24fdd17 100644
--- a/third_party/xla/xla/backends/autotuner/codegen_orchestrator.cc
+++ b/third_party/xla/xla/backends/autotuner/codegen_orchestrator.cc
@@ -143,13 +143,6 @@ CodegenOrchestrator::GetDefaultConfig(const HloInstruction& instr) const {
 
 absl::StatusOr<std::unique_ptr<Executable>> CodegenOrchestrator::Compile(
     const HloInstruction& instr, const Config& config) const {
-  if (options_.exclude_cublas_config &&
-      (config.codegen_backend->backend() ==
-           autotuner::Backend::CUBLASLT_FISSION ||
-       config.codegen_backend->backend() ==
-           autotuner::Backend::HIPBLASLT_FISSION)) {
-    return absl::CancelledError("exclude_cublas_config is set.");
-  }
   VLOG(4) << "Compiling config " << config.ToString() << " for HLO "
           << instr.ToString();
   absl::StatusOr<std::unique_ptr<Executable>> executable =
diff --git a/third_party/xla/xla/backends/autotuner/codegen_orchestrator.h b/third_party/xla/xla/backends/autotuner/codegen_orchestrator.h
index 833d703c4fe120..14108ca2ada57b 100644
--- a/third_party/xla/xla/backends/autotuner/codegen_orchestrator.h
+++ b/third_party/xla/xla/backends/autotuner/codegen_orchestrator.h
@@ -44,9 +44,6 @@ class CodegenOrchestrator {
     std::function<bool(const HloInstruction&, autotuner::Backend)>
         allow_reg_spills_fn =
             [](const HloInstruction&, autotuner::Backend) { return false; };
-    // TODO(b/519059655): Generalize and move to tuner.
-    // If true, do not allow compilation of cublas or rocblas configs.
-    bool exclude_cublas_config = false;
   };
 
   // TODO(b/444398084): Unify Cache::Config and CodegenOrchestrator::Config
diff --git a/third_party/xla/xla/backends/autotuner/codegen_orchestrator_test.cc b/third_party/xla/xla/backends/autotuner/codegen_orchestrator_test.cc
index 2bbdb05f03f372..225d2ce9b67d66 100644
--- a/third_party/xla/xla/backends/autotuner/codegen_orchestrator_test.cc
+++ b/third_party/xla/xla/backends/autotuner/codegen_orchestrator_test.cc
@@ -193,37 +193,6 @@ TEST_F(CodegenOrchestratorTest, GetDefaultConfigFailsWhenNoBackendProvides) {
               StatusIs(absl::StatusCode::kNotFound));
 }
 
-class CodegenOrchestratorParamTest
-    : public CodegenOrchestratorTest,
-      public ::testing::WithParamInterface<autotuner::Backend> {};
-
-TEST_P(CodegenOrchestratorParamTest, ExcludeCublasConfig) {
-  CodegenOrchestrator::Options options;
-  options.exclude_cublas_config = true;
-
-  auto backend = std::make_unique<MockCodegenBackend>();
-  EXPECT_CALL(*backend, backend()).WillRepeatedly(Return(GetParam()));
-  EXPECT_CALL(*backend, name()).WillRepeatedly(Return("mock_backend"));
-  EXPECT_CALL(*backend, Compile(_, _)).Times(0);
-
-  CodegenOrchestrator::Config config{backend.get(),
-                                     GetTestConfig("test_config_1")};
-
-  std::vector<std::unique_ptr<CodegenBackend>> backends;
-  backends.push_back(std::move(backend));
-
-  ASSERT_OK_AND_ASSIGN(auto orchestrator, CodegenOrchestrator::Create(
-                                              std::move(backends), options));
-
-  auto dummy_instr = HloInstruction::CreateConstant(LiteralUtil::CreateR0(1));
-  EXPECT_THAT(orchestrator->Compile(*dummy_instr, config),
-              StatusIs(absl::StatusCode::kCancelled));
-}
-
-INSTANTIATE_TEST_SUITE_P(
-    ExcludeCublasConfigs, CodegenOrchestratorParamTest,
-    ::testing::Values(autotuner::Backend::CUBLASLT_FISSION,
-                      autotuner::Backend::HIPBLASLT_FISSION));
 
 TEST_F(CodegenOrchestratorTest, ConfigsWithRegisterSpillingAreAllowed) {
   CodegenOrchestrator::Options options;
diff --git a/third_party/xla/xla/backends/autotuner/config_selector.cc b/third_party/xla/xla/backends/autotuner/config_selector.cc
index 2862d72781c691..6e1e1a9f5b7e02 100644
--- a/third_party/xla/xla/backends/autotuner/config_selector.cc
+++ b/third_party/xla/xla/backends/autotuner/config_selector.cc
@@ -21,25 +21,40 @@ limitations under the License.
 #include <utility>
 #include <vector>
 
+#include "absl/algorithm/container.h"
 #include "absl/log/log.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_join.h"
 #include "absl/time/time.h"
+#include "absl/types/span.h"
+#include "xla/backends/autotuner/backends.pb.h"
 #include "xla/backends/autotuner/config_runner.h"
 
 namespace xla {
 
 absl::StatusOr<ConfigRunner::ConfigProfile> PickBestConfig(
     std::vector<ConfigRunner::ConfigProfile>& results,
-    int scratch_bytes_window_size_us) {
+    int scratch_bytes_window_size_us,
+    absl::Span<const autotuner::Backend> excluded_backends) {
+  auto is_excluded = [&](const ConfigRunner::ConfigProfile& result) {
+    return result.config.codegen_backend != nullptr &&
+           absl::c_linear_search(excluded_backends,
+                                 result.config.codegen_backend->backend());
+  };
+
   absl::Duration min_duration = absl::InfiniteDuration();
   ConfigRunner::ConfigProfile* best_result = nullptr;
   std::vector<std::string> failures;
   for (ConfigRunner::ConfigProfile& result : results) {
     if (result.failure.has_value()) {
       failures.push_back(result.failure->ToString());
+    } else if (is_excluded(result)) {
+      failures.push_back(absl::StrCat(
+          result.config.ToString(), ": Backend excluded from selection (",
+          autotuner::Backend_Name(result.config.codegen_backend->backend()),
+          ")"));
     } else if (result.duration < min_duration) {
       min_duration = result.duration;
       best_result = &result;
@@ -47,7 +62,8 @@ absl::StatusOr<ConfigRunner::ConfigProfile> PickBestConfig(
   }
 
   if (best_result == nullptr) {
-    std::string message = "All configs failed during profiling.";
+    std::string message =
+        "All configs failed during profiling or were excluded from selection.";
     if (!failures.empty()) {
       absl::StrAppend(&message, "\nFailures (", failures.size(), "):\n",
                       absl::StrJoin(failures, "\n"));
@@ -62,7 +78,8 @@ absl::StatusOr<ConfigRunner::ConfigProfile> PickBestConfig(
   absl::Duration min_duration_with_optimized_scratch_bytes =
       absl::InfiniteDuration();
   for (ConfigRunner::ConfigProfile& result : results) {
-    if (!result.failure.has_value() && result.duration <= duration_limit) {
+    if (!result.failure.has_value() && !is_excluded(result) &&
+        result.duration <= duration_limit) {
       bool current_result_is_better =
           result.scratch_bytes < min_scratch_bytes ||
           (result.scratch_bytes == min_scratch_bytes &&
diff --git a/third_party/xla/xla/backends/autotuner/config_selector.h b/third_party/xla/xla/backends/autotuner/config_selector.h
index 6e890b76b6c7d5..5a6dcaa0b93579 100644
--- a/third_party/xla/xla/backends/autotuner/config_selector.h
+++ b/third_party/xla/xla/backends/autotuner/config_selector.h
@@ -19,13 +19,16 @@ limitations under the License.
 #include <vector>
 
 #include "absl/status/statusor.h"
+#include "absl/types/span.h"
+#include "xla/backends/autotuner/backends.pb.h"
 #include "xla/backends/autotuner/config_runner.h"
 
 namespace xla {
 
 absl::StatusOr<ConfigRunner::ConfigProfile> PickBestConfig(
     std::vector<ConfigRunner::ConfigProfile>& results,
-    int scratch_bytes_window_size_us);
+    int scratch_bytes_window_size_us,
+    absl::Span<const autotuner::Backend> excluded_backends = {});
 
 }  // namespace xla
 
diff --git a/third_party/xla/xla/backends/autotuner/config_selector_test.cc b/third_party/xla/xla/backends/autotuner/config_selector_test.cc
index ef61b12e6b571f..b4563644108e49 100644
--- a/third_party/xla/xla/backends/autotuner/config_selector_test.cc
+++ b/third_party/xla/xla/backends/autotuner/config_selector_test.cc
@@ -135,5 +135,82 @@ TEST(ConfigSelectorTest, IgnoresScratchBytesOutsideWindow) {
   EXPECT_EQ(best.scratch_bytes, 200);
 }
 
+TEST(ConfigSelectorTest, PicksNextBestConfigWhenFastestBackendExcluded) {
+  MockCodegenBackend cublas_backend;
+  EXPECT_CALL(cublas_backend, backend())
+      .WillRepeatedly(::testing::Return(autotuner::Backend::CUBLASLT_FISSION));
+  MockCodegenBackend triton_backend;
+  EXPECT_CALL(triton_backend, backend())
+      .WillRepeatedly(::testing::Return(autotuner::Backend::TRITON));
+
+  std::vector<ConfigRunner::ConfigProfile> profiles;
+  profiles.push_back(CreateProfile(&cublas_backend, "cublas_fast_config",
+                                   absl::Microseconds(10)));
+  profiles.push_back(CreateProfile(&triton_backend, "triton_config_1",
+                                   absl::Microseconds(20)));
+  profiles.push_back(CreateProfile(&triton_backend, "triton_config_2",
+                                   absl::Microseconds(30)));
+
+  std::vector<autotuner::Backend> excluded_backends = {
+      autotuner::Backend::CUBLASLT_FISSION};
+
+  ASSERT_OK_AND_ASSIGN(
+      auto best, PickBestConfig(profiles, /*scratch_bytes_window_size_us=*/0,
+                                excluded_backends));
+  EXPECT_THAT(*best.config.backend_config, ConfigMatcher("triton_config_1"));
+  EXPECT_EQ(best.duration, absl::Microseconds(20));
+}
+
+TEST(ConfigSelectorTest, FailsWhenAllSuccessfulConfigsAreFromExcludedBackends) {
+  MockCodegenBackend cublas_backend;
+  EXPECT_CALL(cublas_backend, backend())
+      .WillRepeatedly(::testing::Return(autotuner::Backend::CUBLASLT_FISSION));
+  MockCodegenBackend triton_backend;
+  EXPECT_CALL(triton_backend, backend())
+      .WillRepeatedly(::testing::Return(autotuner::Backend::TRITON));
+
+  std::vector<ConfigRunner::ConfigProfile> profiles;
+  profiles.push_back(
+      CreateProfile(&cublas_backend, "cublas_config", absl::Microseconds(10)));
+  profiles.push_back(CreateProfile(
+      &triton_backend, "triton_config", absl::Microseconds(20),
+      /*scratch_bytes=*/0,
+      ConfigRunner::Failure{ConfigRunner::FailureKind::kWrongResults,
+                            "wrong results"}));
+
+  std::vector<autotuner::Backend> excluded_backends = {
+      autotuner::Backend::CUBLASLT_FISSION};
+
+  EXPECT_THAT(PickBestConfig(profiles, /*scratch_bytes_window_size_us=*/0,
+                             excluded_backends),
+              StatusIs(absl::StatusCode::kNotFound));
+}
+
+TEST(ConfigSelectorTest, ScratchBytesOptimizationIgnoresExcludedBackends) {
+  MockCodegenBackend cublas_backend;
+  EXPECT_CALL(cublas_backend, backend())
+      .WillRepeatedly(::testing::Return(autotuner::Backend::CUBLASLT_FISSION));
+  MockCodegenBackend triton_backend;
+  EXPECT_CALL(triton_backend, backend())
+      .WillRepeatedly(::testing::Return(autotuner::Backend::TRITON));
+
+  std::vector<ConfigRunner::ConfigProfile> profiles;
+  profiles.push_back(CreateProfile(&cublas_backend, "cublas_least_scratch",
+                                   absl::Microseconds(21), 0));
+  profiles.push_back(CreateProfile(&triton_backend, "triton_fast",
+                                   absl::Microseconds(20), 200));
+  profiles.push_back(CreateProfile(&triton_backend, "triton_opt_scratch",
+                                   absl::Microseconds(25), 100));
+
+  std::vector<autotuner::Backend> excluded_backends = {
+      autotuner::Backend::CUBLASLT_FISSION};
+
+  ASSERT_OK_AND_ASSIGN(
+      auto best, PickBestConfig(profiles, /*scratch_bytes_window_size_us=*/8,
+                                excluded_backends));
+  EXPECT_THAT(*best.config.backend_config, ConfigMatcher("triton_opt_scratch"));
+  EXPECT_EQ(best.scratch_bytes, 100);
+}
+
 }  // namespace
 }  // namespace xla
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/BUILD b/third_party/xla/xla/backends/gpu/codegen/emitters/BUILD
index cb4ac01198a0f5..241c0f37aad122 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/BUILD
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/BUILD
@@ -57,6 +57,26 @@ cc_library(
     ],
 )
 
+xla_cc_test(
+    name = "concatenate_test",
+    srcs = ["concatenate_test.cc"],
+    tags = ["gpu"],
+    deps = [
+        ":concatenate",
+        ":mlir_kernel_emitter",
+        "//xla:debug_options_flags",
+        "//xla:xla_proto_cc",
+        "//xla/hlo/ir:hlo",
+        "//xla/hlo/testlib:hlo_hardware_independent_test_base",
+        "//xla/service/gpu:gpu_device_info_for_tests",
+        "//xla/service/gpu:hlo_fusion_analysis",
+        "//xla/stream_executor:device_description",
+        "//xla/stream_executor:semantic_version",
+        "//xla/tests:xla_internal_test_main",
+        "@com_google_googletest//:gtest",
+    ],
+)
+
 cc_library(
     name = "mlir_kernel_emitter",
     srcs = ["mlir_kernel_emitter.cc"],
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/concatenate.h b/third_party/xla/xla/backends/gpu/codegen/emitters/concatenate.h
index 8322004e9c122f..6e98f46a1df20c 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/concatenate.h
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/concatenate.h
@@ -45,6 +45,7 @@ class ConcatenateFusion final : public MlirKernelEmitter {
   explicit ConcatenateFusion(const HloFusionAnalysis& analysis);
 
   LaunchDimensions launch_dimensions() const override;
+  int unroll_factor() const override { return unroll_factor_; }
 
   std::optional<IndexingMap> ComputeThreadIdToOutputIndexing(
       int64_t root_index, mlir::MLIRContext* ctx) const override;
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/concatenate_test.cc b/third_party/xla/xla/backends/gpu/codegen/emitters/concatenate_test.cc
new file mode 100644
index 00000000000000..1eeedb54473b6a
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/concatenate_test.cc
@@ -0,0 +1,62 @@
+/* Copyright 2026 The OpenXLA Authors. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * ============================================================================
+ */
+#include "xla/backends/gpu/codegen/emitters/concatenate.h"
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "xla/backends/gpu/codegen/emitters/mlir_kernel_emitter.h"
+#include "xla/debug_options_flags.h"
+#include "xla/hlo/ir/hlo_instruction.h"
+#include "xla/hlo/testlib/hlo_hardware_independent_test_base.h"
+#include "xla/service/gpu/gpu_device_info_for_tests.h"
+#include "xla/service/gpu/hlo_fusion_analysis.h"
+#include "xla/stream_executor/device_description.h"
+#include "xla/stream_executor/semantic_version.h"
+#include "xla/xla.pb.h"
+
+namespace xla::gpu {
+namespace {
+
+class ConcatenateFusionTest : public HloHardwareIndependentTestBase {
+ protected:
+  DebugOptions GetDebugOptionsForTest() const override {
+    auto debug_options = GetDebugOptionsFromFlags();
+    debug_options.set_xla_gpu_experimental_max_unroll_factor(32);
+    return debug_options;
+  }
+};
+
+TEST_F(ConcatenateFusionTest, PropagatesUnrollFactorToCompilationPipeline) {
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
+    ENTRY main {
+      p0 = bf16[1024000] parameter(0)
+      p1 = bf16[1024000] parameter(1)
+      ROOT result = bf16[2048000] concatenate(p0, p1), dimensions={0}
+    })"));
+
+  se::DeviceDescription device_info = TestGpuDeviceInfo::B200SXMDeviceInfo();
+  device_info.set_compile_time_toolkit_version(se::SemanticVersion(12, 9, 0));
+  HloFusionAnalysis analysis = HloFusionAnalysis::Create(
+      *module->entry_computation()->root_instruction(), device_info);
+  ConcatenateFusion concatenate_fusion(analysis);
+  const MlirKernelEmitter& compilation_pipeline_emitter = concatenate_fusion;
+
+  // Blackwell with CUDA 12.9 vectorizes up to 256 bits, or 16 BF16 elements.
+  EXPECT_EQ(compilation_pipeline_emitter.unroll_factor(), 16);
+}
+
+}  // namespace
+}  // namespace xla::gpu
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/BUILD b/third_party/xla/xla/backends/gpu/codegen/triton/BUILD
index b0a66259272ec4..34cf8bfd3ca081 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/BUILD
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/BUILD
@@ -365,7 +365,6 @@ xla_test(
         "//xla:xla_proto_cc",
         "//xla/backends/gpu/tests:gpu_pjrt_codegen_test",
         "//xla/backends/gpu/transforms:convert_triton_gemm_config",
-        "//xla/backends/gpu/transforms:hoist_fused_bitcasts",
         "//xla/codegen/xtile:block_level_parameters",
         "//xla/hlo/analysis:symbolic_map",
         "//xla/hlo/ir:hlo",
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/triton_gemm_fusion_test.cc b/third_party/xla/xla/backends/gpu/codegen/triton/triton_gemm_fusion_test.cc
index f92c64043f70a4..5ffa216b50edda 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/triton_gemm_fusion_test.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/triton_gemm_fusion_test.cc
@@ -38,7 +38,6 @@ limitations under the License.
 #include "xla/backends/gpu/codegen/triton/xtile_compiler.h"
 #include "xla/backends/gpu/tests/gpu_pjrt_codegen_test.h"
 #include "xla/backends/gpu/transforms/convert_triton_gemm_config.h"
-#include "xla/backends/gpu/transforms/hoist_fused_bitcasts.h"
 #include "xla/codegen/xtile/block_level_parameters.h"
 #include "xla/error_spec.h"
 #include "xla/hlo/analysis/symbolic_expr.h"
@@ -112,7 +111,6 @@ class TritonTest : public HloInterpreterReferenceMixin<GpuPjRtCodegenTest> {
   GetModuleAndNestedFusionMetadata(absl::string_view hlo_text) {
     ABSL_ASSIGN_OR_RETURN(std::unique_ptr<VerifiedHloModule> module,
                      ParseAndReturnVerifiedModule(hlo_text));
-    ABSL_RETURN_IF_ERROR(HoistFusedBitcasts().Run(module.get()).status());
     ABSL_ASSIGN_OR_RETURN(bool converted, ConvertTritonGemmConfig(
                                          device_description(), &mlir_context_)
                                          .Run(module.get()));
@@ -522,14 +520,14 @@ ENTRY entry {
           tsl::error::RESOURCE_EXHAUSTED,
           ::testing::HasSubstr("Shared memory size limit exceeded")));
 
-  TF_ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata module2_and_metadata,
-                          GetModuleAndNestedFusionMetadata(absl::Substitute(
-                              kHloTextTemplate, 64, 128, 128, 1)));
+  ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata module2_and_metadata,
+                       GetModuleAndNestedFusionMetadata(absl::Substitute(
+                           kHloTextTemplate, 64, 128, 128, 1)));
 
   const HloFusionInstruction* fusion2 = Cast<HloFusionInstruction>(
       module2_and_metadata.computation->FusionInstruction());
 
-  TF_ASSERT_OK_AND_ASSIGN(
+  ASSERT_OK_AND_ASSIGN(
       const auto result,
       TritonWrapper("test_fn", *fusion2, se::GpuComputeCapability{cc},
                     device_info, module2_and_metadata.block_level_parameters,
@@ -577,8 +575,8 @@ ENTRY e {
   DebugOptions debug_options = config.debug_options();
   debug_options.clear_xla_disable_hlo_passes();
   config.set_debug_options(debug_options);
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
-                          GetOptimizedModule(kHloText, config));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       GetOptimizedModule(kHloText, config));
   const HloInstruction* instr = module->entry_computation()->root_instruction();
   EXPECT_THAT(
       instr,
@@ -614,8 +612,8 @@ ENTRY e {
       "num_ctas":1}}}
 })";
 
-  TF_ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata module_and_metadata,
-                          GetModuleAndNestedFusionMetadata(kHloText));
+  ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata module_and_metadata,
+                       GetModuleAndNestedFusionMetadata(kHloText));
   EXPECT_TRUE(
       RunAndCompareNoHloPasses(module_and_metadata.module->ToString(),
                                ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3}));
@@ -899,9 +897,9 @@ ENTRY entry {
       "num_ctas":1}}}
 })";
 
-  TF_ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata module1_and_metadata,
-                          GetModuleAndNestedFusionMetadata(absl::Substitute(
-                              kHloTextTemplate, 512, 512, 32)));
+  ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata module1_and_metadata,
+                       GetModuleAndNestedFusionMetadata(
+                           absl::Substitute(kHloTextTemplate, 512, 512, 32)));
 
   const HloFusionInstruction* fusion1 = Cast<HloFusionInstruction>(
       module1_and_metadata.computation->FusionInstruction());
@@ -913,9 +911,9 @@ ENTRY entry {
                              "Tiling complexity heuristic exceeded"));
 
   // Succeeds if the tiling is not too complex.
-  TF_ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata module2_and_metadata,
-                          GetModuleAndNestedFusionMetadata(
-                              absl::Substitute(kHloTextTemplate, 32, 32, 32)));
+  ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata module2_and_metadata,
+                       GetModuleAndNestedFusionMetadata(
+                           absl::Substitute(kHloTextTemplate, 32, 32, 32)));
 
   const HloFusionInstruction* fusion2 = Cast<HloFusionInstruction>(
       module1_and_metadata.computation->FusionInstruction());
@@ -954,8 +952,8 @@ e {
     "triton_gemm_config": {"block_m":"16","block_n":"16","block_k":"64",
           "num_stages":"1","num_warps":"4","num_ctas":"1"}}}
 })";
-  TF_ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata module_and_metadata,
-                          GetModuleAndNestedFusionMetadata(kHloText));
+  ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata module_and_metadata,
+                       GetModuleAndNestedFusionMetadata(kHloText));
 
   EXPECT_TRUE(
       RunAndCompareNoHloPasses(std::move(module_and_metadata.module),
@@ -987,8 +985,8 @@ e {
           "num_stages":"1","num_warps":"1","num_ctas":"1"}}}
 })";
 
-  TF_ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata module_and_metadata,
-                          GetModuleAndNestedFusionMetadata(kHloText));
+  ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata module_and_metadata,
+                       GetModuleAndNestedFusionMetadata(kHloText));
 
   EXPECT_TRUE(
       RunAndCompareNoHloPasses(std::move(module_and_metadata.module),
@@ -1018,8 +1016,8 @@ ENTRY e {
   ROOT dot = f32[4,5] dot(dot_lhs, dynamic_slice),
           lhs_contracting_dims={0}, rhs_contracting_dims={1}
 })";
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
-                          GetOptimizedModule(kHloText));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       GetOptimizedModule(kHloText));
   EXPECT_THAT(
       module->entry_computation()->root_instruction(),
       GmockMatch(m::Fusion(m::Parameter(), m::Parameter(),
@@ -1061,8 +1059,8 @@ ENTRY e {
              "block_m":"32","block_n":"32","block_k":"32",
              "num_stages":"1","num_warps":"4","num_ctas":"1"}}}
 })";
-  TF_ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata module_and_metadata,
-                          GetModuleAndNestedFusionMetadata(kHloText));
+  ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata module_and_metadata,
+                       GetModuleAndNestedFusionMetadata(kHloText));
   EXPECT_TRUE(
       RunAndCompareNoHloPasses(module_and_metadata.module->ToString(),
                                ErrorSpec{/*aabs=*/1e-4, /*arel=*/1e-6}));
@@ -1108,8 +1106,8 @@ ENTRY e {
              "num_stages":"1","num_warps":"4","num_ctas":"1"}}}
 })",
                                                 GetParam());
-  TF_ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata module_and_metadata,
-                          GetModuleAndNestedFusionMetadata(hlo_text));
+  ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata module_and_metadata,
+                       GetModuleAndNestedFusionMetadata(hlo_text));
   EXPECT_TRUE(
       RunAndCompareNoHloPasses(module_and_metadata.module->ToString(),
                                ErrorSpec{/*aabs=*/1e-4, /*arel=*/1e-6}));
@@ -1157,8 +1155,8 @@ ENTRY e {
              "block_m":"32","block_n":"32","block_k":"32",
              "num_stages":"1","num_warps":"4","num_ctas":"1"}}}
 })";
-  TF_ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata module_and_metadata,
-                          GetModuleAndNestedFusionMetadata(kHloText));
+  ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata module_and_metadata,
+                       GetModuleAndNestedFusionMetadata(kHloText));
   EXPECT_TRUE(
       RunAndCompareNoHloPasses(module_and_metadata.module->ToString(),
                                ErrorSpec{/*aabs=*/1e-4, /*arel=*/1e-6}));
@@ -1202,8 +1200,8 @@ ENTRY e {
              "num_stages":"1","num_warps":"4","num_ctas":"1"}}}
 })";
 
-  TF_ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata module_and_metadata,
-                          GetModuleAndNestedFusionMetadata(kHloText));
+  ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata module_and_metadata,
+                       GetModuleAndNestedFusionMetadata(kHloText));
   EXPECT_TRUE(
       RunAndCompareNoHloPasses(module_and_metadata.module->ToString(),
                                ErrorSpec{/*aabs=*/1e-4, /*arel=*/1e-6}));
@@ -1246,8 +1244,8 @@ ENTRY e {
              "block_m":"32","block_n":"32","block_k":"32",
              "num_stages":"1","num_warps":"4","num_ctas":"1"}}}
 })";
-  TF_ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata module_and_metadata,
-                          GetModuleAndNestedFusionMetadata(kHloText));
+  ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata module_and_metadata,
+                       GetModuleAndNestedFusionMetadata(kHloText));
   EXPECT_TRUE(
       RunAndCompareNoHloPasses(module_and_metadata.module->ToString(),
                                ErrorSpec{/*aabs=*/1e-4, /*arel=*/1e-6}));
@@ -1300,8 +1298,8 @@ ENTRY e {
     lhs_contracting_dims={1}, rhs_contracting_dims={0}
 })";
 
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
-                          GetOptimizedModule(kHloText));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       GetOptimizedModule(kHloText));
 
   EXPECT_THAT(
       module->entry_computation()->root_instruction(),
@@ -1351,8 +1349,8 @@ ENTRY e {
     lhs_contracting_dims={0}, rhs_contracting_dims={0}
 })";
 
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
-                          GetOptimizedModule(kHloText));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       GetOptimizedModule(kHloText));
   EXPECT_THAT(
       module->entry_computation()->root_instruction(),
       GmockMatch(m::Fusion(m::Parameter(), m::Parameter(), m::Parameter())
@@ -1373,8 +1371,8 @@ ENTRY e {
     lhs_contracting_dims={0}, rhs_contracting_dims={1}
 })";
 
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
-                          GetOptimizedModule(kHloText));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       GetOptimizedModule(kHloText));
 
   EXPECT_THAT(
       module->entry_computation()->root_instruction(),
@@ -1403,8 +1401,8 @@ ENTRY e {
   ROOT a = f32[7,16] add(d0, d1)
 })";
 
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
-                          GetOptimizedModule(kHloText));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       GetOptimizedModule(kHloText));
 
   EXPECT_THAT(
       module->entry_computation()->root_instruction(),
@@ -1431,8 +1429,8 @@ ENTRY e {
     lhs_contracting_dims={1}, rhs_contracting_dims={0}
 })";
 
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
-                          GetOptimizedModule(kHloText));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       GetOptimizedModule(kHloText));
   EXPECT_THAT(
       module->entry_computation()->root_instruction(),
       GmockMatch(m::Fusion(m::Parameter(), m::Parameter())
@@ -1454,8 +1452,8 @@ ENTRY e {
     lhs_contracting_dims={1}, rhs_contracting_dims={0}
 })";
 
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
-                          GetOptimizedModule(kHloText));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       GetOptimizedModule(kHloText));
   EXPECT_THAT(
       module->entry_computation()->root_instruction(),
       GmockMatch(m::Fusion(m::Parameter())
@@ -1478,8 +1476,8 @@ ENTRY e {
     lhs_contracting_dims={1}, rhs_contracting_dims={0}
 })";
 
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
-                          GetOptimizedModule(kHloText));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       GetOptimizedModule(kHloText));
   EXPECT_THAT(
       GetNonBitcastRoot(module->entry_computation()),
       GmockMatch(m::Fusion(m::Parameter(), m::Parameter())
@@ -1500,8 +1498,8 @@ ENTRY e {
   ROOT d = f16[60,120] dot(c0, r1),
     lhs_contracting_dims={1}, rhs_contracting_dims={0}
 })";
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
-                          GetOptimizedModule(kHloText));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       GetOptimizedModule(kHloText));
   EXPECT_THAT(
       GetNonBitcastRoot(module->entry_computation()),
       GmockMatch(m::Fusion(m::Parameter(), m::Constant())
@@ -1554,8 +1552,8 @@ e {
     lhs_contracting_dims={1}, rhs_contracting_dims={0}
 })";
 
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
-                          GetOptimizedModule(kHloText));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       GetOptimizedModule(kHloText));
   EXPECT_THAT(
       GetNonBitcastRoot(module->entry_computation()),
       GmockMatch(m::Fusion(m::Parameter(), m::Parameter(), m::Parameter(),
@@ -1581,8 +1579,8 @@ ENTRY e {
   ROOT r = f32[7,16] sine(d)
 })";
 
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
-                          GetOptimizedModule(kHloText));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       GetOptimizedModule(kHloText));
   EXPECT_THAT(
       module->entry_computation()->root_instruction(),
       GmockMatch(m::Fusion(m::Parameter(), m::Parameter())
@@ -1602,8 +1600,8 @@ ENTRY e {
     lhs_contracting_dims={0}, rhs_contracting_dims={1}
 })";
 
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
-                          GetOptimizedModule(kHloText));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       GetOptimizedModule(kHloText));
   EXPECT_THAT(
       module->entry_computation()->root_instruction(),
       GmockMatch(m::Fusion(m::Parameter(), m::Parameter())
@@ -1624,8 +1622,8 @@ ENTRY e {
     lhs_contracting_dims={1}, rhs_contracting_dims={0}
 })";
 
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
-                          GetOptimizedModule(kHloText));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       GetOptimizedModule(kHloText));
   EXPECT_THAT(
       GetNonBitcastRoot(module->entry_computation()),
       GmockMatch(m::Fusion(m::Parameter(), m::Parameter())
@@ -1649,8 +1647,8 @@ ENTRY e {
     lhs_contracting_dims={1}, rhs_contracting_dims={0}
 })";
 
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
-                          GetOptimizedModule(kHloText));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       GetOptimizedModule(kHloText));
   EXPECT_THAT(
       module->entry_computation()->root_instruction(),
       GmockMatch(m::Fusion(m::Parameter(), m::Parameter())
@@ -1673,8 +1671,8 @@ ENTRY e {
     rhs_batch_dims={2}, rhs_contracting_dims={1}
 })";
 
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
-                          GetOptimizedModule(kHloText));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       GetOptimizedModule(kHloText));
   EXPECT_THAT(
       module->entry_computation()->root_instruction(),
       GmockMatch(m::Fusion(m::Parameter(), m::Parameter())
@@ -1724,8 +1722,8 @@ ENTRY e {
   ROOT r = f16[54,22] convert(d)
 })";
 
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
-                          GetOptimizedModule(kHloText));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       GetOptimizedModule(kHloText));
   EXPECT_THAT(
       module->entry_computation()->root_instruction(),
       GmockMatch(m::Fusion(m::Parameter(), m::Parameter())
@@ -1751,8 +1749,8 @@ ENTRY e {
   ROOT r = bf16[350,690]{1,0} multiply(p2, dot.21)
 })";
 
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
-                          GetOptimizedModule(kHloText));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       GetOptimizedModule(kHloText));
   const HloInstruction* instr = module->entry_computation()->root_instruction();
   if (!instr->IsCustomFusion()) {
     instr = instr->operand(0);
@@ -1787,8 +1785,8 @@ ENTRY e {
   ROOT multiply.8808 = bf16[350,690]{1,0} multiply(neg.484, p2)
 })";
 
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
-                          GetOptimizedModule(kHloText));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       GetOptimizedModule(kHloText));
   const HloInstruction* instr = module->entry_computation()->root_instruction();
   if (!instr->IsCustomFusion()) {
     instr = instr->operand(0);
@@ -1822,8 +1820,8 @@ ENTRY e {
   ROOT t1 = bf16[5,42,200,15] transpose(r1), dimensions={0,3,1,2}
 })";
 
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
-                          GetOptimizedModule(kHloText));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       GetOptimizedModule(kHloText));
   const HloInstruction* root = module->entry_computation()->root_instruction();
   EXPECT_THAT(
       root,
@@ -1866,8 +1864,8 @@ ENTRY e {
   ROOT t1 = bf16[5,42,20,150] transpose(r1), dimensions={0,3,1,2}
 })";
 
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
-                          GetOptimizedModule(kHloText));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       GetOptimizedModule(kHloText));
   EXPECT_THAT(
       module->entry_computation()->root_instruction(),
       GmockMatch(m::Fusion(m::Parameter(), m::Parameter())
@@ -1891,8 +1889,8 @@ ENTRY e {
 }
 )";
 
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
-                          GetOptimizedModule(kHloText));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       GetOptimizedModule(kHloText));
 
   EXPECT_THAT(
       module->entry_computation()->root_instruction(),
@@ -1914,8 +1912,8 @@ ENTRY e {
   ROOT dot = f32[1,250000] dot(parameter_0, parameter_1), lhs_batch_dims={0},
     lhs_contracting_dims={1}, rhs_batch_dims={0}, rhs_contracting_dims={1}
 })";
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
-                          GetOptimizedModule(kHloText));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       GetOptimizedModule(kHloText));
 
   EXPECT_THAT(
       GetNonBitcastRoot(module->entry_computation()),
@@ -1937,8 +1935,8 @@ ENTRY e {
   ROOT dot = f32[1,250000] dot(parameter_0, parameter_1), lhs_batch_dims={0},
     lhs_contracting_dims={1}, rhs_batch_dims={0}, rhs_contracting_dims={1}
 })";
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
-                          GetOptimizedModule(kHloText));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       GetOptimizedModule(kHloText));
 
   EXPECT_THAT(
       GetNonBitcastRoot(module->entry_computation()),
@@ -1985,12 +1983,11 @@ ENTRY e {
                          "num_ctas":1}}}
 })";
 
-  TF_ASSERT_OK_AND_ASSIGN(
-      ModuleAndNestedFusionMetadata test_module_and_metadata,
-      GetModuleAndNestedFusionMetadata(kHloText));
+  ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata test_module_and_metadata,
+                       GetModuleAndNestedFusionMetadata(kHloText));
 
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> ref_module,
-                          ParseAndReturnVerifiedModule(kHloTextRef));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> ref_module,
+                       ParseAndReturnVerifiedModule(kHloTextRef));
 
   EXPECT_TRUE(RunAndCompareTwoModules(
       std::move(ref_module), std::move(test_module_and_metadata.module),
@@ -2045,7 +2042,7 @@ ENTRY e {
                          "num_ctas":1}}}
 })";
 
-  TF_ASSERT_OK_AND_ASSIGN(
+  ASSERT_OK_AND_ASSIGN(
       ModuleAndNestedFusionMetadata optin_shmem_module_and_metadata,
       GetModuleAndNestedFusionMetadata(kHloTextOptinShmem));
   const HloFusionInstruction* triton_dot_fusion = Cast<HloFusionInstruction>(
@@ -2054,7 +2051,7 @@ ENTRY e {
   llvm::Triple target_triple(nvptx::TargetTriple());
   std::string data_layout(nvptx::DataLayout());
 
-  TF_ASSERT_OK_AND_ASSIGN(
+  ASSERT_OK_AND_ASSIGN(
       const auto result,
       TritonWrapper("test_fn", *triton_dot_fusion, GpuComputeCapability(),
                     dev_info,
@@ -2091,7 +2088,7 @@ ENTRY e {
                          "num_ctas":1}}}
 })";
 
-  TF_ASSERT_OK_AND_ASSIGN(
+  ASSERT_OK_AND_ASSIGN(
       ModuleAndNestedFusionMetadata low_shmem_module_and_metadata,
       GetModuleAndNestedFusionMetadata(kHloTextLowShmem));
 
@@ -2157,12 +2154,11 @@ ENTRY e {
     calls=loop_fusion
 })";
 
-  TF_ASSERT_OK_AND_ASSIGN(
-      ModuleAndNestedFusionMetadata test_module_and_metadata,
-      GetModuleAndNestedFusionMetadata(kHloTextTest));
+  ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata test_module_and_metadata,
+                       GetModuleAndNestedFusionMetadata(kHloTextTest));
 
-  TF_ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata ref_module_and_metadata,
-                          GetModuleAndNestedFusionMetadata(kHloTextRef));
+  ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata ref_module_and_metadata,
+                       GetModuleAndNestedFusionMetadata(kHloTextRef));
 
   EXPECT_TRUE(
       RunAndCompareTwoModules(std::move(ref_module_and_metadata.module),
@@ -2207,11 +2203,11 @@ ENTRY e {
     backend_config={"gemm_backend_config": {"alpha_real":1,"beta":0,"dot_dimension_numbers":{"lhs_contracting_dimensions":["0"],"rhs_contracting_dimensions":["1"],"lhs_batch_dimensions":[],"rhs_batch_dimensions":[]},"alpha_imag":0,"precision_config":{"operand_precision":["DEFAULT","DEFAULT"]},"epilogue":"DEFAULT"}}
   ROOT get-tuple-element = f32[63,92]{1,0} get-tuple-element((f32[63,92]{1,0}, s8[0]{0}) gemm), index=0
 })";
-  TF_ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata module_and_metadata,
-                          GetModuleAndNestedFusionMetadata(kHloTextTest));
+  ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata module_and_metadata,
+                       GetModuleAndNestedFusionMetadata(kHloTextTest));
 
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> ref_module,
-                          ParseAndReturnVerifiedModule(kHloTextRef));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> ref_module,
+                       ParseAndReturnVerifiedModule(kHloTextRef));
 
   EXPECT_TRUE(RunAndCompareTwoModules(std::move(ref_module),
                                       std::move(module_and_metadata.module),
@@ -2256,11 +2252,11 @@ ENTRY triton_gemm___computation {
   ROOT get-tuple-element = f32[11,45]{1,0} get-tuple-element((f32[11,45]{1,0}, s8[0]{0}) gemm), index=0
 })";
 
-  TF_ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata module_and_metadata,
-                          GetModuleAndNestedFusionMetadata(kHloTextTest));
+  ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata module_and_metadata,
+                       GetModuleAndNestedFusionMetadata(kHloTextTest));
 
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> ref_module,
-                          ParseAndReturnVerifiedModule(kHloTextRef));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> ref_module,
+                       ParseAndReturnVerifiedModule(kHloTextRef));
 
   EXPECT_TRUE(RunAndCompareTwoModules(std::move(ref_module),
                                       std::move(module_and_metadata.module),
@@ -2380,11 +2376,11 @@ ENTRY e {
   ROOT get-tuple-element = f32[32,57]{0,1} get-tuple-element((f32[32,57]{0,1}, s8[0]{0}) gemm), index=0
 })";
 
-  TF_ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata module_and_metadata,
-                          GetModuleAndNestedFusionMetadata(kHloTextTest));
+  ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata module_and_metadata,
+                       GetModuleAndNestedFusionMetadata(kHloTextTest));
 
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> ref_module,
-                          ParseAndReturnVerifiedModule(kHloTextRef));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> ref_module,
+                       ParseAndReturnVerifiedModule(kHloTextRef));
 
   EXPECT_TRUE(RunAndCompareTwoModules(std::move(ref_module),
                                       std::move(module_and_metadata.module),
@@ -2422,8 +2418,8 @@ ENTRY e {
                                           "num_ctas":"1"}}}
 })";
 
-  TF_ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata module_and_metadata,
-                          GetModuleAndNestedFusionMetadata(kHloTextTest));
+  ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata module_and_metadata,
+                       GetModuleAndNestedFusionMetadata(kHloTextTest));
 
   const std::string kHloTextRef = R"(
 HloModule m, is_scheduled=true
@@ -2450,8 +2446,8 @@ ENTRY e {
   ROOT get-tuple-element = bf16[92,63]{1,0} get-tuple-element((bf16[92,63]{1,0}, s8[0]{0}) gemm), index=0
 })";
 
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> ref_module,
-                          ParseAndReturnVerifiedModule(kHloTextRef));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> ref_module,
+                       ParseAndReturnVerifiedModule(kHloTextRef));
 
   EXPECT_TRUE(RunAndCompareTwoModules(std::move(ref_module),
                                       std::move(module_and_metadata.module),
@@ -2496,8 +2492,8 @@ ENTRY e {
                          "num_ctas":"1"}}}
 })";
 
-  TF_ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata module_and_metadata,
-                          GetModuleAndNestedFusionMetadata(kHloTextTest));
+  ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata module_and_metadata,
+                       GetModuleAndNestedFusionMetadata(kHloTextTest));
 
   const std::string kHloTextRef = R"(
 ENTRY e {
@@ -2519,8 +2515,8 @@ ENTRY e {
     lhs_contracting_dims={1}, rhs_contracting_dims={0}
 })";
 
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> ref_module,
-                          ParseAndReturnVerifiedModule(kHloTextRef));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> ref_module,
+                       ParseAndReturnVerifiedModule(kHloTextRef));
 
   EXPECT_TRUE(RunAndCompareTwoModules(std::move(ref_module),
                                       std::move(module_and_metadata.module),
@@ -2533,22 +2529,19 @@ TEST_F(TritonTest, UseTF32For8BitOrLessWithF32) {
 HloModule t
 
 triton_dot {
-  parameter_0 = s32[11,24]{1,0} parameter(0)
-  broadcast = s32[11,24,128]{2,1,0} broadcast(parameter_0),
-  dimensions={0,1}
-  parameter_1 = s32[11,24,128]{2,1,0} parameter(1)
-  compare = pred[11,24,128]{2,1,0} compare(broadcast, parameter_1),
-      direction=EQ
-  bitcast = pred[264,128]{1,0} bitcast(compare)
-  convert = f32[264,128]{1,0} convert(bitcast)
+  parameter_0 = s32[264]{0} parameter(0)
+  broadcast = s32[264,128]{1,0} broadcast(parameter_0), dimensions={0}
+  parameter_1 = s32[264,128]{1,0} parameter(1)
+  compare = pred[264,128]{1,0} compare(broadcast, parameter_1), direction=EQ
+  convert = f32[264,128]{1,0} convert(compare)
   parameter_2 = f32[128,8]{1,0} parameter(2)
   ROOT dot = f32[264,8]{1,0} dot(convert, parameter_2),
       lhs_contracting_dims={1}, rhs_contracting_dims={0}
 }
 
 ENTRY e {
-  p0 = s32[11,24]{1,0} parameter(0)
-  p1 = s32[11,24,128]{2,1,0} parameter(1)
+  p0 = s32[264]{0} parameter(0)
+  p1 = s32[264,128]{1,0} parameter(1)
   p2 = f32[128,8]{1,0} parameter(2)
   ROOT _ = f32[264,8] fusion(p0, p1, p2), kind=kCustom, calls=triton_dot,
     backend_config={"fusion_backend_config": {kind: "__triton_gemm",
@@ -2557,9 +2550,9 @@ ENTRY e {
          "num_stages":1,"num_warps":4,
          "num_ctas":1}}}
 })";
-  TF_ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata module_and_metadata,
-                          GetModuleAndNestedFusionMetadata(kHloText));
-  TF_ASSERT_OK(
+  ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata module_and_metadata,
+                       GetModuleAndNestedFusionMetadata(kHloText));
+  ASSERT_OK(
       CreateTritonIrAndFileCheck(*module_and_metadata.computation,
                                  module_and_metadata.block_level_parameters,
                                  R"(
@@ -2594,8 +2587,8 @@ ENTRY e {
                          "num_ctas":1}}}
 })";
 
-  TF_ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata module_and_metadata,
-                          GetModuleAndNestedFusionMetadata(kHloText));
+  ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata module_and_metadata,
+                       GetModuleAndNestedFusionMetadata(kHloText));
   EXPECT_TRUE(Run(std::move(module_and_metadata.module),
                   /*run_hlo_passes=*/false));
 }
@@ -2622,8 +2615,8 @@ ENTRY e {
                          "num_ctas":1}}}
 })";
 
-  TF_ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata module_and_metadata,
-                          GetModuleAndNestedFusionMetadata(kHloText));
+  ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata module_and_metadata,
+                       GetModuleAndNestedFusionMetadata(kHloText));
   EXPECT_TRUE(Run(std::move(module_and_metadata.module),
                   /*run_hlo_passes=*/false));
 }
@@ -2653,8 +2646,8 @@ ENTRY e {
                          "num_ctas":1}}}
 })";
 
-  TF_ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata module_and_metadata,
-                          GetModuleAndNestedFusionMetadata(kHloText));
+  ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata module_and_metadata,
+                       GetModuleAndNestedFusionMetadata(kHloText));
   EXPECT_TRUE(Run(std::move(module_and_metadata.module),
                   /*run_hlo_passes=*/false));
 }
@@ -2664,8 +2657,8 @@ TEST_F(TritonGemmTest, MixedF8DotExecutesCorrectly) {
     GTEST_SKIP() << "Requires a Hopper+ GPU";
   }
 
-  TF_ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata module_and_metadata,
-                          GetModuleAndNestedFusionMetadata(R"(
+  ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata module_and_metadata,
+                       GetModuleAndNestedFusionMetadata(R"(
 triton_dot {
   p0 = f8e5m2[32,32] parameter(0)
   p1 = f8e4m3fn[32,32] parameter(1)
@@ -2683,8 +2676,8 @@ e {
                          "num_ctas":1}}}
 })"));
 
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> ref_module,
-                          ParseAndReturnVerifiedModule(R"(
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> ref_module,
+                       ParseAndReturnVerifiedModule(R"(
 e {
   p0 = f8e5m2[32,32] parameter(0)
   p0c = f16[32,32] convert(p0)
@@ -2725,8 +2718,8 @@ ENTRY e {
                          "num_ctas":1}}}
 })";
 
-  TF_ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata module_and_metadata,
-                          GetModuleAndNestedFusionMetadata(kHloText));
+  ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata module_and_metadata,
+                       GetModuleAndNestedFusionMetadata(kHloText));
   EXPECT_TRUE(Run(std::move(module_and_metadata.module),
                   /*run_hlo_passes=*/false));
 }
@@ -2757,8 +2750,8 @@ ENTRY e {
          "num_ctas":1}}}
 })";
 
-  TF_ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata module_and_metadata,
-                          GetModuleAndNestedFusionMetadata(kHloText));
+  ASSERT_OK_AND_ASSIGN(ModuleAndNestedFusionMetadata module_and_metadata,
+                       GetModuleAndNestedFusionMetadata(kHloText));
 
   CompileAndOptionallyVerifyPtx(std::move(module_and_metadata.module), R"(
 CHECK: wgmma.mma_async.sync.aligned.m64n16k16.f32.bf16.bf16
@@ -2780,8 +2773,8 @@ ENTRY e {
   ROOT _ = f16[30,30] dot(p0, cp1),
     lhs_contracting_dims={0}, rhs_contracting_dims={1}
 })";
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> verified_module,
-                          ParseAndReturnVerifiedModule(kHloText));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> verified_module,
+                       ParseAndReturnVerifiedModule(kHloText));
   DebugOptions debug_options = verified_module->config().debug_options();
   debug_options.set_xla_gpu_autotune_level(0);
   verified_module->mutable_config().set_debug_options(debug_options);
diff --git a/third_party/xla/xla/backends/gpu/runtime/BUILD b/third_party/xla/xla/backends/gpu/runtime/BUILD
index 56f9a5697b38ce..75bdcc99a3ef38 100644
--- a/third_party/xla/xla/backends/gpu/runtime/BUILD
+++ b/third_party/xla/xla/backends/gpu/runtime/BUILD
@@ -1961,8 +1961,9 @@ cc_library(
     hdrs = ["collective_kernel_thunk.h"],
     deps = [
         ":all_reduce",
-        ":collective_kernel_api",
+        ":collective_cliques",
         ":collective_kernel_thunk_proto_cc",
+        ":collective_memory",
         ":collective_params",
         ":collective_thunk",
         ":collective_thunk_proto_cc",
@@ -1974,7 +1975,10 @@ cc_library(
         "//xla:util",
         "//xla:xla_data_proto_cc",
         "//xla/backends/gpu/collectives:gpu_clique_key",
+        "//xla/backends/gpu/collectives:gpu_clique_rendezvous",
+        "//xla/backends/gpu/collectives:gpu_communicator",
         "//xla/core/collectives:rank_id",
+        "//xla/core/collectives:symmetric_memory",
         "//xla/runtime:buffer_use",
         "//xla/runtime:device_id",
         "//xla/service:buffer_assignment",
@@ -1992,6 +1996,7 @@ cc_library(
         "//xla/stream_executor/gpu:all_reduce_kernel",
         "//xla/stream_executor/gpu:collective_kernel_metadata",
         "//xla/tsl/util:safe_reinterpret_cast",
+        "//xla/tsl/util:tied_ref",
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:flat_hash_map",
@@ -2029,12 +2034,14 @@ xla_test(
         ":collective_thunk",
         ":command_state",
         ":thunk",
+        "//xla:debug_options_flags",
         "//xla:future",
         "//xla:shape_util",
         "//xla:util",
         "//xla:xla_data_proto_cc",
         "//xla:xla_proto_cc",
         "//xla/backends/gpu/collectives:gpu_clique_key",
+        "//xla/backends/gpu/transforms/collectives:collective_ops_utils",
         "//xla/runtime:buffer_use",
         "//xla/runtime:device_id",
         "//xla/service:buffer_assignment",
@@ -3430,8 +3437,6 @@ cc_library(
     compatible_with = get_compatible_with_portable(),
     deps = [
         "//xla:status_macros",
-        "//xla/backends/gpu/collectives:gpu_clique_key",
-        "//xla/backends/gpu/collectives:gpu_clique_rendezvous",
         "//xla/core/collectives:rank_id",
         "//xla/core/collectives:symmetric_memory",
         "//xla/stream_executor:device_address",
diff --git a/third_party/xla/xla/backends/gpu/runtime/all_reduce.cc b/third_party/xla/xla/backends/gpu/runtime/all_reduce.cc
index 4034d62cad5a8d..8a4faab4fa0de0 100644
--- a/third_party/xla/xla/backends/gpu/runtime/all_reduce.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/all_reduce.cc
@@ -449,6 +449,13 @@ absl::StatusOr<CollectiveKernelSpec> CreateAllReduceKernelSpec(
   const int64_t remote_size =
       xla::RoundUpTo<uint64_t>(input_size_bytes, kXlaAllocatedBufferAlignBytes);
 
+  const DebugOptions& debug_options =
+      instr->GetModule()->config().debug_options();
+  const SymmetricMemoryType sym_mem_type =
+      IsCrossHostOneShotKernelEnabled(debug_options, DebugOptions::ALLREDUCE)
+          ? SymmetricMemoryType::kLoadStoreAccessible
+          : SymmetricMemoryType::kXlaRendezvous;
+
   CollectiveKernelSpec kernel_spec = {
       /* .input_buffer_specs= */ {
           {/*requires_multimem=*/false, SymmetricMemoryType::kNone}},
@@ -456,11 +463,11 @@ absl::StatusOr<CollectiveKernelSpec> CreateAllReduceKernelSpec(
       {{/*requires_multimem=*/false, SymmetricMemoryType::kNone}},
       /* .scratch_buffers= */
       {{signal_size, /*requires_multimem=*/false,  // Signal buffers
-        SymmetricMemoryType::kXlaRendezvous,
+        sym_mem_type,
         /*should_memzero=*/true,
         /*should_double_buffer=*/true},
        {remote_size, /*requires_multimem=*/false,  // Remote buffers
-        SymmetricMemoryType::kXlaRendezvous,
+        sym_mem_type,
         /*should_memzero=*/false,
         /*should_double_buffer=*/true}},
       /* .argument_descriptors= */
diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_api.cc b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_api.cc
index dce95619458e2c..0eb6811f600d7b 100644
--- a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_api.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_api.cc
@@ -26,11 +26,8 @@ limitations under the License.
 #include "absl/base/no_destructor.h"
 #include "absl/container/flat_hash_map.h"
 #include "absl/status/status.h"
-#include "absl/status/status_macros.h"
 #include "absl/status/statusor.h"
 #include "absl/synchronization/mutex.h"
-#include "xla/backends/gpu/collectives/gpu_clique_key.h"
-#include "xla/backends/gpu/collectives/gpu_clique_rendezvous.h"
 #include "xla/core/collectives/rank_id.h"
 #include "xla/core/collectives/symmetric_memory.h"
 #include "xla/status_macros.h"
@@ -44,6 +41,7 @@ limitations under the License.
 
 namespace xla {
 namespace gpu {
+namespace se = ::stream_executor;
 
 namespace {
 
@@ -148,41 +146,5 @@ size_t GetMultiGpuBarrierSignalBufferSize() {
 
 size_t GetMultiGpuBarrierSignalValueSize() { return sizeof(uint32_t); }
 
-absl::StatusOr<std::vector<void*>> CollectParamToPeers(
-    const GpuCliqueKey& clique_key, RankId rank,
-    stream_executor::Stream* stream,
-    std::vector<stream_executor::DeviceAddressBase> parameters) {
-  std::vector<void*> param_to_peers_ptrs;
-
-  size_t num_parameters = parameters.size();
-  // Exchange device parameters with all ranks in the clique.
-  ABSL_ASSIGN_OR_RETURN(
-      auto device_parameters,
-      GpuCliqueRendezvous::Join(clique_key, rank, std::move(parameters)));
-
-  // Collect pointers to device buffers from all participating ranks.
-  param_to_peers_ptrs.reserve(num_parameters * clique_key.num_devices());
-
-  absl::flat_hash_map<int, std::vector<stream_executor::DeviceAddressBase>>
-      peer_to_parameters(clique_key.num_devices());
-
-  using DeviceParameters = std::vector<stream_executor::DeviceAddressBase>;
-
-  for (auto peer = RankId(0); peer < RankId(clique_key.num_devices()); ++peer) {
-    ABSL_ASSIGN_OR_RETURN(const DeviceParameters& peer_parameters,
-                     device_parameters->at<DeviceParameters>(peer));
-    peer_to_parameters[peer.value()] = std::move(peer_parameters);
-  }
-
-  for (int parameter = 0; parameter < num_parameters; ++parameter) {
-    for (int peer = 0; peer < clique_key.num_devices(); ++peer) {
-      param_to_peers_ptrs.push_back(
-          peer_to_parameters[peer][parameter].opaque());
-    }
-  }
-
-  return param_to_peers_ptrs;
-}
-
 }  // namespace gpu
 }  // namespace xla
diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_api.h b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_api.h
index ff4bc5101d7593..cd601d10610dd9 100644
--- a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_api.h
+++ b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_api.h
@@ -21,8 +21,6 @@ limitations under the License.
 #include <vector>
 
 #include "absl/status/status.h"
-#include "absl/status/statusor.h"
-#include "xla/backends/gpu/collectives/gpu_clique_key.h"
 #include "xla/core/collectives/rank_id.h"
 #include "xla/core/collectives/symmetric_memory.h"
 #include "xla/stream_executor/device_address.h"
@@ -57,13 +55,6 @@ size_t GetMultiGpuBarrierSignalBufferSize();
 // Returns the size of the barrier signal value in bytes.
 size_t GetMultiGpuBarrierSignalValueSize();
 
-// Collect the pointers to the parameters at the peer devices.
-// The size of the returned vector is num_parameters * num_devices.
-absl::StatusOr<std::vector<void*>> CollectParamToPeers(
-    const GpuCliqueKey& clique_key, RankId rank,
-    stream_executor::Stream* stream,
-    std::vector<stream_executor::DeviceAddressBase> parameters);
-
 }  // namespace xla::gpu
 
 #endif  // XLA_BACKENDS_GPU_RUNTIME_COLLECTIVE_KERNEL_API_H_
diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.cc
index 5294bcb0bc783c..a14258fc9d372c 100644
--- a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.cc
@@ -23,6 +23,7 @@ limitations under the License.*/
 #include <vector>
 
 #include "absl/algorithm/container.h"
+#include "absl/container/flat_hash_map.h"
 #include "absl/log/log.h"
 #include "absl/log/vlog_is_on.h"
 #include "absl/status/status.h"
@@ -35,15 +36,19 @@ limitations under the License.*/
 #include "absl/synchronization/mutex.h"
 #include "absl/types/span.h"
 #include "xla/backends/gpu/collectives/gpu_clique_key.h"
+#include "xla/backends/gpu/collectives/gpu_clique_rendezvous.h"
+#include "xla/backends/gpu/collectives/gpu_communicator.h"
 #include "xla/backends/gpu/runtime/all_reduce.h"
-#include "xla/backends/gpu/runtime/collective_kernel_api.h"
+#include "xla/backends/gpu/runtime/collective_cliques.h"
 #include "xla/backends/gpu/runtime/collective_kernel_thunk.pb.h"
+#include "xla/backends/gpu/runtime/collective_memory.h"
 #include "xla/backends/gpu/runtime/collective_params.h"
 #include "xla/backends/gpu/runtime/collective_thunk.h"
 #include "xla/backends/gpu/runtime/collective_thunk.pb.h"
 #include "xla/backends/gpu/runtime/thunk.h"
 #include "xla/backends/gpu/runtime/thunk.pb.h"
 #include "xla/core/collectives/rank_id.h"
+#include "xla/core/collectives/symmetric_memory.h"
 #include "xla/runtime/buffer_use.h"
 #include "xla/runtime/device_id.h"
 #include "xla/service/buffer_assignment.h"
@@ -62,6 +67,7 @@ limitations under the License.*/
 #include "xla/stream_executor/stream.h"
 #include "xla/stream_executor/stream_executor.h"
 #include "xla/tsl/util/safe_reinterpret_cast.h"
+#include "xla/tsl/util/tied_ref.h"
 #include "xla/util.h"
 #include "xla/xla_data.pb.h"
 
@@ -242,6 +248,43 @@ bool RequiresMultimem(CollectiveKernelSpec kernel_spec) {
   }
   return false;
 }
+
+absl::StatusOr<std::vector<void*>> CollectParamToPeers(
+    const GpuCliqueKey& clique_key, RankId rank,
+    stream_executor::Stream* stream,
+    std::vector<stream_executor::DeviceAddressBase> parameters) {
+  std::vector<void*> param_to_peers_ptrs;
+
+  size_t num_parameters = parameters.size();
+  // Exchange device parameters with all ranks in the clique.
+  ABSL_ASSIGN_OR_RETURN(
+      auto device_parameters,
+      GpuCliqueRendezvous::Join(clique_key, rank, std::move(parameters)));
+
+  // Collect pointers to device buffers from all participating ranks.
+  param_to_peers_ptrs.reserve(num_parameters * clique_key.num_devices());
+
+  absl::flat_hash_map<int, std::vector<stream_executor::DeviceAddressBase>>
+      peer_to_parameters(clique_key.num_devices());
+
+  using DeviceParameters = std::vector<stream_executor::DeviceAddressBase>;
+
+  for (auto peer = RankId(0); peer < RankId(clique_key.num_devices()); ++peer) {
+    ABSL_ASSIGN_OR_RETURN(const DeviceParameters& peer_parameters,
+                     device_parameters->at<DeviceParameters>(peer));
+    peer_to_parameters[peer.value()] = std::move(peer_parameters);
+  }
+
+  for (int parameter = 0; parameter < num_parameters; ++parameter) {
+    for (int peer = 0; peer < clique_key.num_devices(); ++peer) {
+      param_to_peers_ptrs.push_back(
+          peer_to_parameters[peer][parameter].opaque());
+    }
+  }
+
+  return param_to_peers_ptrs;
+}
+
 }  // namespace
 
 absl::Status CollectiveKernelThunk::IsSupported(
@@ -455,8 +498,45 @@ absl::Status CollectiveKernelThunk::Initialize(const InitializeParams& params) {
     const size_t num_parameters = parameters.size();
     const size_t param_to_peers_ptrs_size_bytes =
         num_parameters * clique_key.num_devices() * sizeof(uint64_t);
+    TF_RET_CHECK(params.collective_params != nullptr)
+        << "Collective params must not be null in "
+           "CollectiveKernelThunk::Initialize";
+    TF_RET_CHECK(params.collective_cliques != nullptr)
+        << "Collective cliques must not be null in "
+           "CollectiveKernelThunk::Initialize";
+    TF_RET_CHECK(params.collective_memory != nullptr)
+        << "Collective memory must not be null in "
+           "CollectiveKernelThunk::Initialize";
+
+    const bool use_symmetric_memory = absl::c_any_of(
+        kernel_spec_.scratch_buffers, [](const ScratchBufferSpec& spec) {
+          return spec.symmetric_memory_type ==
+                 SymmetricMemoryType::kLoadStoreAccessible;
+        });
+
+    if (use_symmetric_memory) {
+      ABSL_ASSIGN_OR_RETURN(GpuCommunicator * comm,
+                       params.collective_cliques->GetComm(clique_key, *rank));
+
+      if (memory_state->scratch_symmetric_memories.empty()) {
+        memory_state->scratch_symmetric_memories.reserve(
+            memory_state->scratch_allocations.size());
+        for (size_t i = 0; i < memory_state->scratch_allocations.size(); ++i) {
+          se::DeviceAddressBase addr =
+              memory_state->scratch_allocations[i].address();
+          ABSL_ASSIGN_OR_RETURN(std::unique_ptr<SymmetricMemory> symmetric_memory,
+                           comm->CreateSymmetricMemory(addr));
+          ABSL_ASSIGN_OR_RETURN(tsl::TiedRef<SymmetricMemory> tied_symmetric_memory,
+                           params.collective_cliques->Tie(
+                               clique_key, std::move(symmetric_memory)));
+          memory_state->scratch_symmetric_memories.push_back(
+              std::move(tied_symmetric_memory));
+        }
+      }
+    }
+
     std::vector<void*> multimem_addresses;
-    if (RequiresMultimem(kernel_spec_) && params.collective_memory != nullptr) {
+    if (RequiresMultimem(kernel_spec_)) {
       multimem_addresses.resize(num_parameters, nullptr);
       for (size_t i = 0; i < num_parameters; ++i) {
         auto [mmem, offset] = params.collective_memory->FindSymmetricMemory(
@@ -469,9 +549,41 @@ absl::Status CollectiveKernelThunk::Initialize(const InitializeParams& params) {
         }
       }
     }
-    ABSL_ASSIGN_OR_RETURN(std::vector<void*> param_to_peers_ptrs,
-                     CollectParamToPeers(clique_key, state->rank, params.stream,
-                                         std::move(parameters)));
+
+    std::vector<void*> param_to_peers_ptrs;
+    if (use_symmetric_memory) {
+      static constexpr auto is_multimem_buffer =
+          [](const IoBufferSpec& spec) -> bool {
+        return spec.requires_multimem;
+      };
+      int32_t scratch_buffers_index =
+          absl::c_count_if(kernel_spec_.input_buffer_specs,
+                           is_multimem_buffer) +
+          absl::c_count_if(kernel_spec_.output_buffer_specs,
+                           is_multimem_buffer);
+      param_to_peers_ptrs.resize(num_parameters * clique_key.num_devices());
+      for (size_t i = scratch_buffers_index; i < num_parameters; ++i) {
+        const size_t scratch_index = i - scratch_buffers_index;
+        auto sym_mem =
+            memory_state->scratch_symmetric_memories[scratch_index].Lock();
+        TF_RET_CHECK(sym_mem != nullptr)
+            << "Symmetric memory for scratch buffer " << scratch_index
+            << " is no longer valid";
+        const size_t parameter_offset = i * clique_key.num_devices();
+        for (int device_rank = 0; device_rank < clique_key.num_devices();
+             ++device_rank) {
+          ABSL_ASSIGN_OR_RETURN(se::DeviceAddressBase peer_address,
+                           sym_mem->peer_addr(RankId(device_rank)));
+          param_to_peers_ptrs[parameter_offset + device_rank] =
+              peer_address.opaque();
+        }
+      }
+    } else {
+      ABSL_ASSIGN_OR_RETURN(
+          param_to_peers_ptrs,
+          CollectParamToPeers(clique_key, state->rank, params.stream,
+                              std::move(parameters)));
+    }
     const size_t multimem_size_bytes =
         multimem_addresses.size() * sizeof(void*);
     state->metadata = params.executor->Allocate(
diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.h b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.h
index 0316d660cf4b81..760b3643accba4 100644
--- a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.h
+++ b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.h
@@ -37,6 +37,7 @@ limitations under the License.*/
 #include "xla/backends/gpu/runtime/thunk.pb.h"
 #include "xla/backends/gpu/runtime/traced_command.h"
 #include "xla/core/collectives/rank_id.h"
+#include "xla/core/collectives/symmetric_memory.h"
 #include "xla/service/buffer_assignment.h"
 #include "xla/service/gpu/launch_dimensions.h"
 #include "xla/stream_executor/device_address.h"
@@ -44,6 +45,7 @@ limitations under the License.*/
 #include "xla/stream_executor/gpu/all_reduce_kernel.h"
 #include "xla/stream_executor/kernel.h"
 #include "xla/stream_executor/stream.h"
+#include "xla/tsl/util/tied_ref.h"
 
 namespace xla::gpu {
 
@@ -132,6 +134,7 @@ class CollectiveKernelThunk : public TracedCommand {
   // Per-executor scratch memory.
   struct StreamMemory {
     std::vector<se::DeviceAddressHandle> scratch_allocations;
+    std::vector<tsl::TiedRef<SymmetricMemory>> scratch_symmetric_memories;
   };
 
   // Per-executor state that needs to be synchronized for access.
diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk_test.cc b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk_test.cc
index ec65e9093d7f75..0270630937bd49 100644
--- a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk_test.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk_test.cc
@@ -41,6 +41,8 @@ limitations under the License.
 #include "xla/backends/gpu/runtime/collective_thunk.h"
 #include "xla/backends/gpu/runtime/command_state.h"
 #include "xla/backends/gpu/runtime/thunk.h"
+#include "xla/backends/gpu/transforms/collectives/collective_ops_utils.h"
+#include "xla/debug_options_flags.h"
 #include "xla/future.h"
 #include "xla/runtime/buffer_use.h"
 #include "xla/runtime/device_id.h"
@@ -168,23 +170,33 @@ struct CollectiveKernelThunkMetadata {
   std::vector<CollectiveThunk::Buffer> buffers;
 };
 
-CollectiveKernelSpec CreateCollectiveKernelSpec(int64_t num_elements,
-                                                int64_t signal_size,
-                                                int64_t remote_size,
-                                                bool is_multimem_enabled) {
+DebugOptions DefaultDebugOptions() {
+  DebugOptions debug_options = DefaultDebugOptionsIgnoringFlags();
+  debug_options.add_xla_gpu_unsupported_use_cross_host_one_shot_kernel(
+      DebugOptions::ALLCOLLECTIVES);
+  return debug_options;
+}
+
+CollectiveKernelSpec CreateCollectiveKernelSpec(
+    int64_t num_elements, int64_t signal_size, int64_t remote_size,
+    bool is_multimem_enabled,
+    const DebugOptions& debug_options = DefaultDebugOptions()) {
+  const SymmetricMemoryType sym_mem_type =
+      IsCrossHostOneShotKernelEnabled(debug_options,
+                                      DebugOptions::ALLCOLLECTIVES)
+          ? SymmetricMemoryType::kLoadStoreAccessible
+          : SymmetricMemoryType::kXlaRendezvous;
   return {
       /*operand_buffer_specs=*/{
           {/*requires_multimem=*/false, SymmetricMemoryType::kNone}},
       /*result_buffer_specs=*/
       {{/*requires_multimem=*/false, SymmetricMemoryType::kNone}},
       /*scratch_buffers=*/
-      {{signal_size, /*requires_multimem=*/false,
-        SymmetricMemoryType::kXlaRendezvous,
+      {{signal_size, /*requires_multimem=*/false, sym_mem_type,
         /*should_memzero=*/true,
         /*should_double_buffer=*/true},
        {remote_size,
-        /*requires_multimem=*/is_multimem_enabled,
-        SymmetricMemoryType::kXlaRendezvous,
+        /*requires_multimem=*/is_multimem_enabled, sym_mem_type,
         /*should_memzero=*/false,
         /*should_double_buffer=*/true}},
       /*argument_descriptors=*/
@@ -198,7 +210,8 @@ CollectiveKernelSpec CreateCollectiveKernelSpec(int64_t num_elements,
 }
 
 CollectiveKernelThunkMetadata CreateCollectiveKernelThunk(
-    int num_devices, int num_elements, bool is_multimem_enabled, bool use_ptx) {
+    int num_devices, int num_elements, bool is_multimem_enabled, bool use_ptx,
+    const DebugOptions& debug_options = DefaultDebugOptions()) {
   const int64_t input_size_bytes = num_elements * sizeof(uint64_t);
   Shape input_shape = ShapeUtil::MakeShape(U64, {num_elements});
   ReplicaGroup replica_group;
@@ -239,7 +252,7 @@ CollectiveKernelThunkMetadata CreateCollectiveKernelThunk(
   result.thunk = std::make_unique<CollectiveKernelThunk>(
       std::move(thunk_info), collective_config,
       CreateCollectiveKernelSpec(num_elements, signal_size, remote_size,
-                                 is_multimem_enabled),
+                                 is_multimem_enabled, debug_options),
       result.buffers, /*is_collective_kernel_enabled=*/true,
       /*kernel_name=*/kKernelName,
       /*launch_dimensions=*/launch_dimensions,
@@ -372,6 +385,7 @@ absl::StatusOr<se::DeviceAddressBase> RunCollectiveKernelThunk(
   initialize_params.stream = stream.get();
   initialize_params.buffer_allocations = &buffer_allocations;
   initialize_params.collective_params = &collective_params;
+  initialize_params.collective_cliques = &collective_cliques;
   initialize_params.src = {kKernelSource};
   initialize_params.collective_memory = &collective_memory;
 
@@ -474,7 +488,8 @@ TEST(CollectiveKernelThunkTest, MultiprocessTest) {
 
   CollectiveKernelThunkMetadata metadata = CreateCollectiveKernelThunk(
       /*num_devices=*/kDevicesCount, /*num_elements=*/kNumElements,
-      /*is_multimem_enabled=*/false, /*use_ptx=*/true);
+      /*is_multimem_enabled=*/false, /*use_ptx=*/true,
+      /*debug_options=*/DebugOptions());
   EXPECT_THAT(RunCollectiveKernelThunkOnDevices(metadata,
                                                 /*emulate_multiprocess=*/true),
               StatusIs(absl::StatusCode::kInvalidArgument));
@@ -615,19 +630,29 @@ TEST(CollectiveKernelThunkTest, RecordCommandBufferCreateUpdate) {
                                       &allocations1};
   ASSERT_OK(collective_kernel_thunk->Prepare(prepare_params));
 
+  CollectiveMemoryCache collective_memory_cache;
+  ASSERT_OK_AND_ASSIGN(
+      CollectiveCliques collective_cliques,
+      AcquireCollectiveCliques(collective_params, clique_requests));
+  ASSERT_OK_AND_ASSIGN(
+      CollectiveMemory collective_memory,
+      AcquireCollectiveMemory(collective_params, collective_cliques,
+                              memory_requests, collective_memory_cache));
+
   Thunk::InitializeParams initialize_params;
   initialize_params.executor = executor;
   initialize_params.stream = stream.get();
   initialize_params.buffer_allocations = &allocations1;
   initialize_params.collective_params = &collective_params;
+  initialize_params.collective_cliques = &collective_cliques;
   initialize_params.src.text = kKernelSource;
+  initialize_params.collective_memory = &collective_memory;
   ASSERT_OK(collective_kernel_thunk->Initialize(initialize_params));
   ASSERT_OK(stream->BlockHostUntilDone());
 
   Thunk::ExecuteParams params1 = Thunk::ExecuteParams::Create(
       run_options, allocations1, stream.get(), trace_stream.get(),
-      &collective_params, /*collective_cliques=*/nullptr,
-      /*collective_memory=*/nullptr);
+      &collective_params, &collective_cliques, &collective_memory);
 
   CommandStateManager state;
   Command::RecordParams record_params = {state};
@@ -648,8 +673,7 @@ TEST(CollectiveKernelThunkTest, RecordCommandBufferCreateUpdate) {
   BufferAllocations updated_allocations({src2, dst2}, 0, nullptr);
   Thunk::ExecuteParams params2 = Thunk::ExecuteParams::Create(
       run_options, updated_allocations, stream.get(), trace_stream.get(),
-      &collective_params, /*collective_cliques=*/nullptr,
-      /*collective_memory=*/nullptr);
+      &collective_params, &collective_cliques, &collective_memory);
   std::vector<BufferAllocation::Index> updated_allocs = {0, 1};
   Command::RecordParams update_record_params = {state,
                                                 std::move(updated_allocs)};
diff --git a/third_party/xla/xla/backends/gpu/transforms/BUILD b/third_party/xla/xla/backends/gpu/transforms/BUILD
index 58c39a4071ce6d..3a1d75ccf57470 100644
--- a/third_party/xla/xla/backends/gpu/transforms/BUILD
+++ b/third_party/xla/xla/backends/gpu/transforms/BUILD
@@ -45,6 +45,7 @@ xla_cc_test(
     deps = [
         ":dot_algorithm_rewriter",
         "//xla:xla_data_proto_cc",
+        "//xla:xla_proto_cc",
         "//xla/hlo/testlib:filecheck",
         "//xla/hlo/testlib:hlo_hardware_independent_test_base",
         "//xla/service:hlo_module_config",
@@ -109,6 +110,7 @@ cc_library(
         "//xla:literal_util",
         "//xla:shape_util",
         "//xla:util",
+        "//xla:xla_data_proto_cc",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/pass:hlo_pass",
         "@com_google_absl//absl/container:flat_hash_set",
@@ -527,6 +529,7 @@ cc_library(
         "//xla/hlo/ir:hlo",
         "//xla/hlo/pass:hlo_pass",
         "//xla/service:hlo_creation_utils",
+        "//xla/service:hlo_proto_cc",
         "//xla/service:shape_inference",
         "//xla/service/gpu:conv_utils",
         "//xla/service/gpu:cublas_cudnn",
@@ -571,6 +574,7 @@ cc_library(
         "//xla:xla_data_proto_cc",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/pass:hlo_pass",
+        "//xla/service:hlo_proto_cc",
         "//xla/service/gpu:backend_configs_cc",
         "//xla/service/gpu:cublas_cudnn",
         "//xla/stream_executor:device_description",
@@ -667,6 +671,7 @@ cc_library(
         "//xla:xla_data_proto_cc",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/pass:hlo_pass",
+        "//xla/service:hlo_proto_cc",
         "//xla/service/gpu:backend_configs_cc",
         "//xla/service/gpu:conv_utils",
         "//xla/stream_executor:device_description",
@@ -731,6 +736,7 @@ xla_cc_test(
         "//xla/hlo/testlib:pattern_matcher_gmock",
         "//xla/hlo/testlib:test",
         "//xla/hlo/testlib:test_helpers",
+        "//xla/service:hlo_proto_cc",
         "//xla/service:pattern_matcher",
         "//xla/service:shape_inference",
         "//xla/stream_executor:device_description",
@@ -823,6 +829,7 @@ cc_library(
         "//xla/hlo/analysis:hlo_reachability",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/pass:hlo_pass",
+        "//xla/service:hlo_proto_cc",
         "//xla/service/gpu:backend_configs_cc",
         "//xla/service/gpu:ir_emission_utils",
         "//xla/stream_executor:device_description",
@@ -852,6 +859,8 @@ xla_test(
         ":conv_fusion_rewriter",
         ":conv_kind_assignment",
         "//xla:error_spec",
+        "//xla:xla_data_proto_cc",
+        "//xla:xla_proto_cc",
         "//xla/backends/gpu/tests:hlo_pjrt_gpu_test_base",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/testlib:filecheck",
@@ -985,6 +994,7 @@ cc_library(
         "@com_google_absl//absl/container:inlined_vector",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/log:vlog_is_on",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_macros",
         "@com_google_absl//absl/status:statusor",
@@ -1206,6 +1216,7 @@ cc_library(
         "//xla:shape_util",
         "//xla:types",
         "//xla:util",
+        "//xla:xla_data_proto_cc",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/pass:hlo_pass",
         "//xla/service:hlo_creation_utils",
@@ -1234,6 +1245,7 @@ xla_test(
     tags = ["cuda-only"],
     deps = [
         "//xla:error_spec",
+        "//xla:xla_proto_cc",
         "//xla/backends/gpu/tests:hlo_pjrt_gpu_test_base",
         "//xla/stream_executor:device_description",
         "//xla/stream_executor:semantic_version",
@@ -1308,6 +1320,7 @@ cc_library(
         "@com_google_absl//absl/container:inlined_vector",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/log:vlog_is_on",
         "@com_google_absl//absl/status:status_macros",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
@@ -2153,6 +2166,7 @@ xla_test(
         ":gemm_rewriter",
         ":gemm_rewriter_test_lib",
         "//xla:error_spec",
+        "//xla:xla_data_proto_cc",
         "//xla:xla_proto_cc",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/testlib:filecheck",
@@ -2180,6 +2194,7 @@ xla_test(
     backends = ["gpu"],
     use_legacy_runtime = True,
     deps = [
+        "//xla:xla_proto_cc",
         "//xla/backends/gpu/tests:hlo_legacy_gpu_test_base",
         "//xla/hlo/ir:hlo",
         "//xla/service:executable",
@@ -2672,6 +2687,8 @@ xla_test(
     deps = [
         ":ragged_dot_fusion_rewriter",
         "//xla:error_spec",
+        "//xla:xla_data_proto_cc",
+        "//xla:xla_proto_cc",
         "//xla/backends/gpu/tests:hlo_pjrt_gpu_test_base",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/testlib:filecheck",
@@ -3063,7 +3080,9 @@ cc_library(
         "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status:status_macros",
+        "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings:str_format",
+        "@com_google_absl//absl/strings:string_view",
         "@com_google_absl//absl/types:span",
     ],
 )
@@ -3076,6 +3095,7 @@ xla_test(
         ":scatter_determinism_expander",
         "//xla:literal",
         "//xla:shape_util",
+        "//xla:xla_data_proto_cc",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/testlib:test",
         "//xla/tests:hlo_test_base",
@@ -3232,6 +3252,7 @@ xla_cc_test(
     deps = [
         ":softmax_rewriter_triton",
         "//xla:xla_data_proto_cc",
+        "//xla:xla_proto_cc",
         "//xla/backends/gpu/codegen/triton:support",
         "//xla/hlo/analysis:symbolic_map",
         "//xla/hlo/ir:hlo",
@@ -3297,6 +3318,7 @@ cc_library(
     deps = [
         "//xla:shape_util",
         "//xla:status_macros",
+        "//xla:xla_data_proto_cc",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/pass:hlo_pass",
         "//xla/hlo/utils:hlo_query",
@@ -3335,6 +3357,7 @@ cc_library(
         "//xla:status_macros",
         "//xla:util",
         "//xla:xla_data_proto_cc",
+        "//xla:xla_proto_cc",
         "//xla/backends/gpu/libraries/cub:cub_scratch_size_deviceless_lookup",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/pass:hlo_pass",
@@ -4267,6 +4290,7 @@ xla_cc_test(
     deps = [
         ":sort_rewriter",
         "//xla:xla_data_proto_cc",
+        "//xla:xla_proto_cc",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/testlib:hlo_hardware_independent_test_base",
         "//xla/hlo/testlib:pattern_matcher_gmock",
diff --git a/third_party/xla/xla/backends/gpu/transforms/collectives/BUILD b/third_party/xla/xla/backends/gpu/transforms/collectives/BUILD
index 42d3d593ae886f..a29d92582ca5b0 100644
--- a/third_party/xla/xla/backends/gpu/transforms/collectives/BUILD
+++ b/third_party/xla/xla/backends/gpu/transforms/collectives/BUILD
@@ -204,10 +204,12 @@ cc_library(
     deps = [
         "//xla:side_effect_util",
         "//xla:xla_data_proto_cc",
+        "//xla:xla_proto_cc",
         "//xla/hlo/ir:hlo",
         "//xla/runtime:device_id",
         "//xla/service:collective_ops_utils",
         "//xla/service:device_assignment",
+        "//xla/service:gpu_topology",
         "//xla/service:hlo_module_config",
         "//xla/service/gpu:backend_configs_cc",
         "//xla/stream_executor:device_description",
@@ -233,6 +235,8 @@ xla_cc_test(
         ":collective_ops_utils",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/parser:hlo_parser",
+        "//xla/service:device_assignment",
+        "//xla/service:gpu_topology",
         "//xla/service/gpu:backend_configs_cc",
         "//xla/service/gpu:gpu_device_info_for_tests",
         "//xla/stream_executor:device_description",
diff --git a/third_party/xla/xla/backends/gpu/transforms/collectives/collective_ops_utils.cc b/third_party/xla/xla/backends/gpu/transforms/collectives/collective_ops_utils.cc
index 46d98160522060..1838bfe46ff948 100644
--- a/third_party/xla/xla/backends/gpu/transforms/collectives/collective_ops_utils.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/collectives/collective_ops_utils.cc
@@ -331,6 +331,20 @@ bool IsSpmdGenerated(const HloInstruction& instr) {
   return backend_config->collective_backend_config().is_spmd_generated();
 }
 
+bool IsCrossHostOneShotKernelEnabled(
+    const DebugOptions& debug_options,
+    std::optional<DebugOptions::CollectiveOpType> op_type) {
+  if (!op_type.has_value()) {
+    return false;
+  }
+  return absl::c_linear_search(
+             debug_options.xla_gpu_unsupported_use_cross_host_one_shot_kernel(),
+             *op_type) ||
+         absl::c_linear_search(
+             debug_options.xla_gpu_unsupported_use_cross_host_one_shot_kernel(),
+             DebugOptions::ALLCOLLECTIVES);
+}
+
 bool IsAllReplicasLocal(int64_t gpus_per_host,
                         absl::Span<const ReplicaGroup> replica_groups,
                         CollectiveOpGroupMode group_mode,
diff --git a/third_party/xla/xla/backends/gpu/transforms/collectives/collective_ops_utils.h b/third_party/xla/xla/backends/gpu/transforms/collectives/collective_ops_utils.h
index 06b8685c47ac97..b0f2aa38379b70 100644
--- a/third_party/xla/xla/backends/gpu/transforms/collectives/collective_ops_utils.h
+++ b/third_party/xla/xla/backends/gpu/transforms/collectives/collective_ops_utils.h
@@ -88,6 +88,12 @@ bool IsGPUSyncCollective(const HloInstruction& instr);
 // Returns true if all devices are within the same NVLink domain (slice).
 bool IsIntraNVLinkDomain(const HloModuleConfig& config, int64_t slice_size);
 
+// Returns true if xla_gpu_unsupported_use_cross_host_one_shot_kernel is enabled
+// for the given collective op type.
+bool IsCrossHostOneShotKernelEnabled(
+    const DebugOptions& debug_options,
+    std::optional<DebugOptions::CollectiveOpType> op_type);
+
 // Returns true if all replicas in every replica group of the collective
 // are located on the same host (node).
 bool IsAllReplicasLocal(int64_t gpus_per_host,
diff --git a/third_party/xla/xla/backends/gpu/transforms/conv_canonicalizer.cc b/third_party/xla/xla/backends/gpu/transforms/conv_canonicalizer.cc
index 4f89955a73f6ed..a223f2a3acc698 100644
--- a/third_party/xla/xla/backends/gpu/transforms/conv_canonicalizer.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/conv_canonicalizer.cc
@@ -32,6 +32,7 @@ limitations under the License.
 #include "xla/shape.h"
 #include "xla/shape_util.h"
 #include "xla/util.h"
+#include "xla/xla_data.pb.h"
 
 namespace xla {
 namespace gpu {
diff --git a/third_party/xla/xla/backends/gpu/transforms/conv_fusion_rewriter.cc b/third_party/xla/xla/backends/gpu/transforms/conv_fusion_rewriter.cc
index e6224421ff59a8..7cef77dadcac8a 100644
--- a/third_party/xla/xla/backends/gpu/transforms/conv_fusion_rewriter.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/conv_fusion_rewriter.cc
@@ -43,6 +43,7 @@ limitations under the License.
 #include "xla/primitive_util.h"
 #include "xla/service/gpu/backend_configs.pb.h"
 #include "xla/service/gpu/ir_emission_utils.h"
+#include "xla/service/hlo.pb.h"
 #include "xla/shape.h"
 #include "xla/shape_util.h"
 #include "xla/stream_executor/cuda/cuda_compute_capability.h"
diff --git a/third_party/xla/xla/backends/gpu/transforms/conv_fusion_rewriter_test.cc b/third_party/xla/xla/backends/gpu/transforms/conv_fusion_rewriter_test.cc
index 59a515813eb372..7cf3fa00701406 100644
--- a/third_party/xla/xla/backends/gpu/transforms/conv_fusion_rewriter_test.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/conv_fusion_rewriter_test.cc
@@ -50,6 +50,8 @@ limitations under the License.
 #include "xla/tests/hlo_test_base.h"
 #include "xla/tsl/lib/core/status_test_util.h"
 #include "xla/tsl/platform/statusor.h"
+#include "xla/xla.pb.h"
+#include "xla/xla_data.pb.h"
 
 namespace xla {
 namespace gpu {
diff --git a/third_party/xla/xla/backends/gpu/transforms/conv_kind_assignment.cc b/third_party/xla/xla/backends/gpu/transforms/conv_kind_assignment.cc
index 7a9935bcc0c208..fb88ca15673b2e 100644
--- a/third_party/xla/xla/backends/gpu/transforms/conv_kind_assignment.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/conv_kind_assignment.cc
@@ -39,6 +39,7 @@ limitations under the License.
 #include "xla/primitive_util.h"
 #include "xla/service/gpu/backend_configs.pb.h"
 #include "xla/service/gpu/conv_utils.h"
+#include "xla/service/hlo.pb.h"
 #include "xla/shape.h"
 #include "xla/shape_util.h"
 #include "xla/stream_executor/cuda/cuda_compute_capability.h"
diff --git a/third_party/xla/xla/backends/gpu/transforms/conv_kind_assignment_test.cc b/third_party/xla/xla/backends/gpu/transforms/conv_kind_assignment_test.cc
index 7c4bf0651a8114..69fc44208a866c 100644
--- a/third_party/xla/xla/backends/gpu/transforms/conv_kind_assignment_test.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/conv_kind_assignment_test.cc
@@ -36,6 +36,7 @@ limitations under the License.
 #include "xla/hlo/testlib/test.h"
 #include "xla/hlo/testlib/test_helpers.h"
 #include "xla/literal_util.h"
+#include "xla/service/hlo.pb.h"
 #include "xla/service/pattern_matcher.h"
 #include "xla/service/shape_inference.h"
 #include "xla/shape_util.h"
diff --git a/third_party/xla/xla/backends/gpu/transforms/conv_padding_legalization.cc b/third_party/xla/xla/backends/gpu/transforms/conv_padding_legalization.cc
index 9c2ffeb72d2e75..3d8d09146434e5 100644
--- a/third_party/xla/xla/backends/gpu/transforms/conv_padding_legalization.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/conv_padding_legalization.cc
@@ -15,7 +15,6 @@ limitations under the License.
 
 #include "xla/backends/gpu/transforms/conv_padding_legalization.h"
 
-#include <algorithm>
 #include <cstddef>
 #include <cstdint>
 #include <cstdlib>
@@ -36,6 +35,7 @@ limitations under the License.
 #include "xla/literal_util.h"
 #include "xla/service/gpu/conv_utils.h"
 #include "xla/service/gpu/cublas_cudnn.h"
+#include "xla/service/hlo.pb.h"
 #include "xla/service/hlo_creation_utils.h"
 #include "xla/service/shape_inference.h"
 #include "xla/shape.h"
diff --git a/third_party/xla/xla/backends/gpu/transforms/conv_rewriter.cc b/third_party/xla/xla/backends/gpu/transforms/conv_rewriter.cc
index 685fb149dc5ba4..db80fe0e2f0504 100644
--- a/third_party/xla/xla/backends/gpu/transforms/conv_rewriter.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/conv_rewriter.cc
@@ -41,6 +41,7 @@ limitations under the License.
 #include "xla/primitive_util.h"
 #include "xla/service/gpu/backend_configs.pb.h"
 #include "xla/service/gpu/cublas_cudnn.h"
+#include "xla/service/hlo.pb.h"
 #include "xla/shape.h"
 #include "xla/shape_util.h"
 #include "xla/stream_executor/cuda/cuda_compute_capability.h"
diff --git a/third_party/xla/xla/backends/gpu/transforms/cudnn_fused_conv_rewriter.cc b/third_party/xla/xla/backends/gpu/transforms/cudnn_fused_conv_rewriter.cc
index 5e0766b1e5da0c..bb6f8587dd4448 100644
--- a/third_party/xla/xla/backends/gpu/transforms/cudnn_fused_conv_rewriter.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/cudnn_fused_conv_rewriter.cc
@@ -33,6 +33,7 @@ limitations under the License.
 #include "absl/container/inlined_vector.h"
 #include "absl/log/check.h"
 #include "absl/log/log.h"
+#include "absl/log/vlog_is_on.h"
 #include "absl/status/status.h"
 #include "absl/status/status_macros.h"
 #include "absl/strings/str_cat.h"
diff --git a/third_party/xla/xla/backends/gpu/transforms/cudnn_norm_rewriter.cc b/third_party/xla/xla/backends/gpu/transforms/cudnn_norm_rewriter.cc
index 97c53ebfb015fe..d664cfa395a2fb 100644
--- a/third_party/xla/xla/backends/gpu/transforms/cudnn_norm_rewriter.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/cudnn_norm_rewriter.cc
@@ -47,6 +47,7 @@ limitations under the License.
 #include "xla/tsl/protobuf/dnn.pb.h"
 #include "xla/types.h"
 #include "xla/util.h"
+#include "xla/xla_data.pb.h"
 
 namespace xla {
 namespace gpu {
diff --git a/third_party/xla/xla/backends/gpu/transforms/cudnn_norm_rewriter_test.cc b/third_party/xla/xla/backends/gpu/transforms/cudnn_norm_rewriter_test.cc
index a90df9c8f451ab..ae22ab531862f4 100644
--- a/third_party/xla/xla/backends/gpu/transforms/cudnn_norm_rewriter_test.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/cudnn_norm_rewriter_test.cc
@@ -23,6 +23,7 @@ limitations under the License.
 #include "xla/stream_executor/device_description.h"
 #include "xla/stream_executor/semantic_version.h"
 #include "xla/tests/hlo_pjrt_interpreter_reference_mixin.h"
+#include "xla/xla.pb.h"
 
 namespace xla::gpu {
 namespace {
diff --git a/third_party/xla/xla/backends/gpu/transforms/cudnn_simplify_padding.cc b/third_party/xla/xla/backends/gpu/transforms/cudnn_simplify_padding.cc
index beb8a8cb0db3df..ca0f04ae538676 100644
--- a/third_party/xla/xla/backends/gpu/transforms/cudnn_simplify_padding.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/cudnn_simplify_padding.cc
@@ -25,6 +25,7 @@ limitations under the License.
 #include "absl/container/inlined_vector.h"
 #include "absl/log/check.h"
 #include "absl/log/log.h"
+#include "absl/log/vlog_is_on.h"
 #include "absl/status/status_macros.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/str_join.h"
diff --git a/third_party/xla/xla/backends/gpu/transforms/dot_algorithm_rewriter_test.cc b/third_party/xla/xla/backends/gpu/transforms/dot_algorithm_rewriter_test.cc
index 53201572720e4d..47d4b3e68276ac 100644
--- a/third_party/xla/xla/backends/gpu/transforms/dot_algorithm_rewriter_test.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/dot_algorithm_rewriter_test.cc
@@ -22,6 +22,7 @@ limitations under the License.
 #include "xla/hlo/testlib/filecheck.h"
 #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h"
 #include "xla/service/hlo_module_config.h"
+#include "xla/xla.pb.h"
 #include "xla/xla_data.pb.h"
 
 namespace xla::gpu {
diff --git a/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_fusion_test.cc b/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_fusion_test.cc
index 3959c4879f76e6..99ddc830851833 100644
--- a/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_fusion_test.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_fusion_test.cc
@@ -16,6 +16,7 @@ limitations under the License.
 #include "xla/backends/gpu/transforms/dynamic_slice_fusion.h"
 
 #include <cstdint>
+#include <optional>
 #include <string>
 #include <vector>
 
diff --git a/third_party/xla/xla/backends/gpu/transforms/gemm_rewriter_allocation_test.cc b/third_party/xla/xla/backends/gpu/transforms/gemm_rewriter_allocation_test.cc
index c4bd0c7b2a31e4..1587c475d12c12 100644
--- a/third_party/xla/xla/backends/gpu/transforms/gemm_rewriter_allocation_test.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/gemm_rewriter_allocation_test.cc
@@ -25,6 +25,7 @@ limitations under the License.
 #include "xla/stream_executor/device_address_allocator.h"
 #include "xla/stream_executor/stream_executor_address_allocator.h"
 #include "xla/tsl/platform/statusor.h"
+#include "xla/xla.pb.h"
 
 namespace xla::gpu {
 namespace {
diff --git a/third_party/xla/xla/backends/gpu/transforms/gemm_rewriter_test.cc b/third_party/xla/xla/backends/gpu/transforms/gemm_rewriter_test.cc
index 966001508cd7c0..87e0143e99bce5 100644
--- a/third_party/xla/xla/backends/gpu/transforms/gemm_rewriter_test.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/gemm_rewriter_test.cc
@@ -46,6 +46,7 @@ limitations under the License.
 #include "xla/tests/hlo_pjrt_interpreter_reference_mixin.h"
 #include "xla/tsl/platform/statusor.h"
 #include "xla/xla.pb.h"
+#include "xla/xla_data.pb.h"
 
 namespace xla {
 namespace gpu {
diff --git a/third_party/xla/xla/backends/gpu/transforms/ragged_all_to_all_multi_host_decomposer.cc b/third_party/xla/xla/backends/gpu/transforms/ragged_all_to_all_multi_host_decomposer.cc
index c36db84bff2fb1..c1cf64ec55e22c 100644
--- a/third_party/xla/xla/backends/gpu/transforms/ragged_all_to_all_multi_host_decomposer.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/ragged_all_to_all_multi_host_decomposer.cc
@@ -16,7 +16,6 @@ limitations under the License.
 #include "xla/backends/gpu/transforms/ragged_all_to_all_multi_host_decomposer.h"
 
 #include <cstdint>
-#include <iterator>
 #include <memory>
 #include <optional>
 #include <vector>
diff --git a/third_party/xla/xla/backends/gpu/transforms/ragged_dot_fusion_rewriter.cc b/third_party/xla/xla/backends/gpu/transforms/ragged_dot_fusion_rewriter.cc
index 12a55a7015a35f..4280a2d98ca083 100644
--- a/third_party/xla/xla/backends/gpu/transforms/ragged_dot_fusion_rewriter.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/ragged_dot_fusion_rewriter.cc
@@ -15,12 +15,10 @@ limitations under the License.
 
 #include "xla/backends/gpu/transforms/ragged_dot_fusion_rewriter.h"
 
-#include <algorithm>
 #include <cstdint>
-#include <iterator>
 #include <memory>
+#include <string>
 #include <utility>
-#include <variant>
 #include <vector>
 
 #include "absl/container/flat_hash_set.h"
@@ -28,6 +26,7 @@ limitations under the License.
 #include "absl/status/status.h"
 #include "absl/status/status_macros.h"
 #include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
 #include "absl/types/span.h"
 #include "llvm/ADT/SmallVector.h"
diff --git a/third_party/xla/xla/backends/gpu/transforms/ragged_dot_fusion_rewriter_test.cc b/third_party/xla/xla/backends/gpu/transforms/ragged_dot_fusion_rewriter_test.cc
index 17caded20a40b4..eef16c9ff276f6 100644
--- a/third_party/xla/xla/backends/gpu/transforms/ragged_dot_fusion_rewriter_test.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/ragged_dot_fusion_rewriter_test.cc
@@ -15,7 +15,6 @@ limitations under the License.
 
 #include "xla/backends/gpu/transforms/ragged_dot_fusion_rewriter.h"
 
-#include <array>
 #include <initializer_list>
 #include <memory>
 #include <string>
@@ -24,7 +23,6 @@ limitations under the License.
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
-#include "absl/container/flat_hash_map.h"
 #include "absl/log/log.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/str_replace.h"
@@ -48,6 +46,8 @@ limitations under the License.
 #include "xla/tests/hlo_pjrt_test_base.h"
 #include "xla/tsl/lib/core/status_test_util.h"
 #include "xla/tsl/platform/statusor.h"
+#include "xla/xla.pb.h"
+#include "xla/xla_data.pb.h"
 
 namespace xla {
 namespace gpu {
diff --git a/third_party/xla/xla/backends/gpu/transforms/scan_rewriter.cc b/third_party/xla/xla/backends/gpu/transforms/scan_rewriter.cc
index 554f120a701594..b47bcf1ca3e1c2 100644
--- a/third_party/xla/xla/backends/gpu/transforms/scan_rewriter.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/scan_rewriter.cc
@@ -34,6 +34,7 @@ limitations under the License.
 #include "xla/shape.h"
 #include "xla/shape_util.h"
 #include "xla/status_macros.h"
+#include "xla/xla_data.pb.h"
 
 namespace xla::gpu {
 
diff --git a/third_party/xla/xla/backends/gpu/transforms/scatter_determinism_expander.h b/third_party/xla/xla/backends/gpu/transforms/scatter_determinism_expander.h
index 91d8d2bd01bad8..d2e7b54d49a055 100644
--- a/third_party/xla/xla/backends/gpu/transforms/scatter_determinism_expander.h
+++ b/third_party/xla/xla/backends/gpu/transforms/scatter_determinism_expander.h
@@ -16,6 +16,8 @@ limitations under the License.
 #ifndef XLA_BACKENDS_GPU_TRANSFORMS_SCATTER_DETERMINISM_EXPANDER_H_
 #define XLA_BACKENDS_GPU_TRANSFORMS_SCATTER_DETERMINISM_EXPANDER_H_
 
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
 #include "xla/hlo/transforms/expanders/op_expander_pass.h"
 
 namespace xla {
diff --git a/third_party/xla/xla/backends/gpu/transforms/scatter_determinism_expander_test.cc b/third_party/xla/xla/backends/gpu/transforms/scatter_determinism_expander_test.cc
index e183feb08fd3d9..935ceb79c7efe5 100644
--- a/third_party/xla/xla/backends/gpu/transforms/scatter_determinism_expander_test.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/scatter_determinism_expander_test.cc
@@ -30,6 +30,7 @@ limitations under the License.
 #include "xla/primitive_util.h"
 #include "xla/shape_util.h"
 #include "xla/tests/hlo_test_base.h"
+#include "xla/xla_data.pb.h"
 
 namespace xla {
 namespace {
diff --git a/third_party/xla/xla/backends/gpu/transforms/softmax_rewriter_triton_test.cc b/third_party/xla/xla/backends/gpu/transforms/softmax_rewriter_triton_test.cc
index de4d464e4e60e2..518e99ab073e3b 100644
--- a/third_party/xla/xla/backends/gpu/transforms/softmax_rewriter_triton_test.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/softmax_rewriter_triton_test.cc
@@ -41,6 +41,7 @@ limitations under the License.
 #include "xla/stream_executor/cuda/cuda_compute_capability.h"
 #include "xla/stream_executor/device_description.h"
 #include "xla/tsl/platform/errors.h"
+#include "xla/xla.pb.h"
 #include "xla/xla_data.pb.h"
 
 namespace xla {
diff --git a/third_party/xla/xla/backends/gpu/transforms/sort_rewriter.cc b/third_party/xla/xla/backends/gpu/transforms/sort_rewriter.cc
index 53127ed9ae78d0..c5779f7b5a7c86 100644
--- a/third_party/xla/xla/backends/gpu/transforms/sort_rewriter.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/sort_rewriter.cc
@@ -54,6 +54,7 @@ limitations under the License.
 #include "xla/tsl/platform/logging.h"
 #include "xla/tsl/platform/statusor.h"
 #include "xla/util.h"
+#include "xla/xla.pb.h"
 #include "xla/xla_data.pb.h"
 
 namespace xla::gpu {
diff --git a/third_party/xla/xla/backends/gpu/transforms/sort_rewriter.h b/third_party/xla/xla/backends/gpu/transforms/sort_rewriter.h
index d7dbcd393e8256..6730ac63610efc 100644
--- a/third_party/xla/xla/backends/gpu/transforms/sort_rewriter.h
+++ b/third_party/xla/xla/backends/gpu/transforms/sort_rewriter.h
@@ -27,6 +27,7 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_module.h"
 #include "xla/hlo/pass/hlo_pass_interface.h"
 #include "xla/stream_executor/device_description.h"
+#include "xla/xla.pb.h"
 
 namespace xla {
 namespace gpu {
diff --git a/third_party/xla/xla/backends/gpu/transforms/sort_rewriter_deviceless_test.cc b/third_party/xla/xla/backends/gpu/transforms/sort_rewriter_deviceless_test.cc
index 1d69f30d705b39..81f39038c359eb 100644
--- a/third_party/xla/xla/backends/gpu/transforms/sort_rewriter_deviceless_test.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/sort_rewriter_deviceless_test.cc
@@ -32,6 +32,7 @@ limitations under the License.
 #include "xla/service/pattern_matcher.h"
 #include "xla/stream_executor/device_description.h"
 #include "xla/stream_executor/semantic_version.h"
+#include "xla/xla.pb.h"
 #include "xla/xla_data.pb.h"
 
 namespace xla {
diff --git a/third_party/xla/xla/backends/gpu/transforms/sort_rewriter_test.cc b/third_party/xla/xla/backends/gpu/transforms/sort_rewriter_test.cc
index f8c3e5bf570c6a..66d5226901ea28 100644
--- a/third_party/xla/xla/backends/gpu/transforms/sort_rewriter_test.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/sort_rewriter_test.cc
@@ -24,7 +24,6 @@ limitations under the License.
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include "absl/status/status_matchers.h"
-#include "absl/strings/ascii.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
 #include "absl/strings/substitute.h"
diff --git a/third_party/xla/xla/codegen/emitters/transforms/expand_integer_power.cc b/third_party/xla/xla/codegen/emitters/transforms/expand_integer_power.cc
index 23678c37ef094a..fce0866c1aff6c 100644
--- a/third_party/xla/xla/codegen/emitters/transforms/expand_integer_power.cc
+++ b/third_party/xla/xla/codegen/emitters/transforms/expand_integer_power.cc
@@ -48,8 +48,9 @@ mlir::LogicalResult ExpandIntegerPower(mlir::math::IPowIOp op,
   llvm::SmallVector<mlir::Type> arg_types(op->getOperandTypes());
   mlir::Value result =
       mlir::mhlo::impl::mapMhloOpToStdScalarOp<mlir::mhlo::PowOp>(
-          op.getLoc(), result_types, arg_types, {op->getOperands()},
-          op->getAttrs(), &rewriter);
+          op.getLoc(), result_types, arg_types,
+          mlir::mhlo::PowOp::Adaptor(op->getOperands()), op->getAttrs(),
+          &rewriter);
 
   rewriter.replaceOp(op, result);
   return mlir::success();
diff --git a/third_party/xla/xla/codegen/xtile/BUILD b/third_party/xla/xla/codegen/xtile/BUILD
index 2501c050badf55..fe521fde604bed 100644
--- a/third_party/xla/xla/codegen/xtile/BUILD
+++ b/third_party/xla/xla/codegen/xtile/BUILD
@@ -74,6 +74,7 @@ xla_cc_test(
         ":tiling_from_block_parameters",
         ":xtile_config_proto_cc",
         "//xla:status_macros",
+        "//xla:xla_proto_cc",
         "//xla/codegen/tiling:symbolic_tile_analysis",
         "//xla/codegen/tiling:tiling_specification",
         "//xla/codegen/tiling/experimental:tile",
diff --git a/third_party/xla/xla/codegen/xtile/tiling_from_block_parameters_test.cc b/third_party/xla/xla/codegen/xtile/tiling_from_block_parameters_test.cc
index 15804e9ff12c6b..cf6c8edb5ca92a 100644
--- a/third_party/xla/xla/codegen/xtile/tiling_from_block_parameters_test.cc
+++ b/third_party/xla/xla/codegen/xtile/tiling_from_block_parameters_test.cc
@@ -45,6 +45,7 @@ limitations under the License.
 #include "xla/service/gpu/backend_configs.pb.h"
 #include "xla/status_macros.h"
 #include "xla/tsl/platform/test.h"
+#include "xla/xla.pb.h"
 
 namespace xla::xtile {
 namespace {
diff --git a/third_party/xla/xla/debug_options_flags.cc b/third_party/xla/xla/debug_options_flags.cc
index 73534629847b27..540880d59af90f 100644
--- a/third_party/xla/xla/debug_options_flags.cc
+++ b/third_party/xla/xla/debug_options_flags.cc
@@ -2233,6 +2233,22 @@ void MakeDebugOptionsFlags(std::vector<tsl::Flag>* flag_list,
       "Only collectives specified in this filter will be executed in a "
       "command buffer. Default is ALLCOLLECTIVES."));
 
+  flag_list->push_back(tsl::Flag(
+      "xla_gpu_unsupported_use_cross_host_one_shot_kernel",
+      SetterForRepeatedEnum<DebugOptions::CollectiveOpType>(
+          "xla_gpu_unsupported_use_cross_host_one_shot_kernel",
+          /*enum_prefix=*/"",
+          [](absl::string_view s, DebugOptions::CollectiveOpType* v) {
+            return DebugOptions::CollectiveOpType_Parse(s, v);
+          },
+          [debug_options]() {
+            return debug_options
+                ->mutable_xla_gpu_unsupported_use_cross_host_one_shot_kernel();
+          }),
+      collective_op_types_to_string(
+          debug_options->xla_gpu_unsupported_use_cross_host_one_shot_kernel()),
+      "Enable cross-host one-shot kernel for specified collectives."));
+
   flag_list->push_back(tsl::Flag(
       "xla_gpu_graph_min_graph_size",
       int32_setter_for(&DebugOptions::set_xla_gpu_graph_min_graph_size),
diff --git a/third_party/xla/xla/debug_options_parsers_test.cc b/third_party/xla/xla/debug_options_parsers_test.cc
index 6f7c08c81084d9..72762e9c7083f3 100644
--- a/third_party/xla/xla/debug_options_parsers_test.cc
+++ b/third_party/xla/xla/debug_options_parsers_test.cc
@@ -210,6 +210,39 @@ TEST(ParsingDebugOptionsTest, ParsingRepeatedFields) {
             DebugOptions::ALLTOALL);
 }
 
+TEST(ParsingDebugOptionsTest, ParsingCrossHostOneShotKernel) {
+  DebugOptions debug_options = DefaultDebugOptionsIgnoringFlags();
+  EXPECT_TRUE(debug_options.xla_gpu_unsupported_use_cross_host_one_shot_kernel()
+                  .empty());
+  debug_options.add_xla_gpu_unsupported_use_cross_host_one_shot_kernel(
+      DebugOptions::ALLGATHER);
+  debug_options.add_xla_gpu_unsupported_use_cross_host_one_shot_kernel(
+      DebugOptions::REDUCESCATTER);
+
+  ResetFlagValues();
+  std::string contents;
+  ASSERT_TRUE(ParseFlagsFromDebugOptionsFile(
+      WriteDebugOptionsToTempFile(debug_options, &contents)));
+  DebugOptions parsed_debug_options = GetDebugOptionsFromFlags();
+  EXPECT_TRUE(absl::StrContains(
+      contents,
+      "xla_gpu_unsupported_use_cross_host_one_shot_kernel: ALLGATHER"));
+  EXPECT_TRUE(absl::StrContains(
+      contents,
+      "xla_gpu_unsupported_use_cross_host_one_shot_kernel: REDUCESCATTER"));
+  EXPECT_EQ(parsed_debug_options
+                .xla_gpu_unsupported_use_cross_host_one_shot_kernel_size(),
+            2);
+  EXPECT_EQ(
+      parsed_debug_options.xla_gpu_unsupported_use_cross_host_one_shot_kernel(
+          0),
+      DebugOptions::ALLGATHER);
+  EXPECT_EQ(
+      parsed_debug_options.xla_gpu_unsupported_use_cross_host_one_shot_kernel(
+          1),
+      DebugOptions::REDUCESCATTER);
+}
+
 TEST(ParsingDebugOptionsTest, ParseFromDebugOptionsFile) {
   // Sanity checks: The test needs to use two flags that have false and true
   // default values.
diff --git a/third_party/xla/xla/hlo/analysis/BUILD b/third_party/xla/xla/hlo/analysis/BUILD
index 325e933827c79a..d0b2e163a1da12 100644
--- a/third_party/xla/xla/hlo/analysis/BUILD
+++ b/third_party/xla/xla/hlo/analysis/BUILD
@@ -465,11 +465,9 @@ cc_library(
         ":alias_info",
         ":hlo_dataflow_analysis",
         ":hlo_operand_index",
-        ":hlo_ordering",
         "//xla:comparison_util",
         "//xla:shape_util",
         "//xla:status_macros",
-        "//xla:types",
         "//xla:util",
         "//xla:xla_data_proto_cc",
         "//xla/hlo/ir:hlo",
@@ -477,8 +475,8 @@ cc_library(
         "//xla/service:hlo_buffer",
         "//xla/service:hlo_value",
         "//xla/tsl/platform:logging",
-        "//xla/tsl/platform:statusor",
         "@com_google_absl//absl/algorithm:container",
+        "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/log",
diff --git a/third_party/xla/xla/hlo/analysis/hlo_alias_analysis.cc b/third_party/xla/xla/hlo/analysis/hlo_alias_analysis.cc
index 6dc3b5aa10f240..856fe28197eb74 100644
--- a/third_party/xla/xla/hlo/analysis/hlo_alias_analysis.cc
+++ b/third_party/xla/xla/hlo/analysis/hlo_alias_analysis.cc
@@ -33,6 +33,10 @@ limitations under the License.
 #include "absl/status/status_macros.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
+#include "absl/strings/str_join.h"
+#include "absl/strings/str_split.h"
+#include "absl/strings/string_view.h"
+#include "absl/types/span.h"
 #include "xla/comparison_util.h"
 #include "xla/hlo/analysis/alias_info.h"
 #include "xla/hlo/analysis/hlo_dataflow_analysis.h"
@@ -48,7 +52,6 @@ limitations under the License.
 #include "xla/shape_util.h"
 #include "xla/status_macros.h"
 #include "xla/tsl/platform/logging.h"
-#include "xla/tsl/platform/statusor.h"
 #include "xla/util.h"
 
 namespace xla {
@@ -332,6 +335,50 @@ std::vector<HloBuffer> CreateBuffers(const HloDataflowAnalysis& dataflow,
 
 HloAliasAnalysis::HloAliasAnalysis(const HloModule* module) : module_(module) {}
 
+bool HloAliasAnalysis::HasLinearCallerChainToRoot(
+    const HloComputation* computation) const {
+  absl::flat_hash_set<const HloComputation*> visited;
+  auto update_has_linear_caller_chain_to_root = [&visited, this](bool value) {
+    for (const HloComputation* visited_computation : visited) {
+      has_linear_caller_chain_to_root_[visited_computation] = value;
+    }
+  };
+
+  while (computation != module_->entry_computation()) {
+    // Check cache.
+    auto it = has_linear_caller_chain_to_root_.find(computation);
+    if (it != has_linear_caller_chain_to_root_.end()) {
+      update_has_linear_caller_chain_to_root(it->second);
+      return it->second;
+    }
+    if (!visited.insert(computation).second ||
+        computation->caller_instructions().size() != 1) {
+      // Recursive computation, or more than one caller, or dead code with zero
+      // callers.
+      update_has_linear_caller_chain_to_root(false);
+      return false;
+    }
+    // Move to the next computation in the caller chain.
+    computation = computation->caller_instructions()[0]->parent();
+  }
+  update_has_linear_caller_chain_to_root(true);
+  return true;
+}
+
+bool HloAliasAnalysis::MultipleBuffersAllowedBeforeInlining(
+    absl::Span<const HloBuffer* const> buffers) const {
+  for (const HloBuffer* buffer : buffers) {
+    for (const HloValue* value : buffer->values()) {
+      for (const HloPosition& position : value->positions()) {
+        if (!HasLinearCallerChainToRoot(position.instruction->parent())) {
+          return true;
+        }
+      }
+    }
+  }
+  return false;
+}
+
 const HloBuffer& HloAliasAnalysis::GetUniqueBufferAt(
     const HloInstruction* instruction, const ShapeIndex& index) const {
   std::vector<const HloBuffer*> buffers = ComputeBuffersAt(instruction, index);
@@ -358,6 +405,27 @@ std::vector<const HloBuffer*> HloAliasAnalysis::ComputeBuffersAt(
   absl::c_sort(buffers, HloBuffer::IdLessThan);
   buffers.erase(std::unique(buffers.begin(), buffers.end()), buffers.end());
 
+  if (buffers.size() > 1 && !MultipleBuffersAllowedBeforeInlining(buffers)) {
+    std::string fingerprint = module_ != nullptr
+                                  ? std::string(module_->GetFingerprint128())
+                                  : "unknown";
+    const bool is_leaf = ShapeUtil::IsLeafIndex(instruction->shape(), index);
+
+    std::string buffers_str = absl::StrJoin(
+        buffers, "\n", [](std::string* out, const HloBuffer* buffer) {
+          absl::StrAppend(out, "    ", buffer->ToString());
+        });
+
+    std::string crash_message =
+        absl::StrCat("More than one buffer found at position:\n",
+                     "  HLO Module Fingerprint: ", fingerprint, "\n",
+                     "  Instruction: ", instruction->name(), "\n",
+                     "  Shape Index: ", index.ToString(), "\n",
+                     "  Is Leaf Index: ", is_leaf ? "true" : "false", "\n",
+                     "  HLO Buffers:\n", buffers_str, "\n");
+    LOG(FATAL) << crash_message;
+  }
+
   return buffers;
 }
 
diff --git a/third_party/xla/xla/hlo/analysis/hlo_alias_analysis.h b/third_party/xla/xla/hlo/analysis/hlo_alias_analysis.h
index 57d06821e48f88..94be322a4ba9dc 100644
--- a/third_party/xla/xla/hlo/analysis/hlo_alias_analysis.h
+++ b/third_party/xla/xla/hlo/analysis/hlo_alias_analysis.h
@@ -21,6 +21,7 @@ limitations under the License.
 #include <vector>
 
 #include "absl/algorithm/container.h"
+#include "absl/base/attributes.h"
 #include "absl/container/flat_hash_map.h"
 #include "absl/container/flat_hash_set.h"
 #include "absl/status/status.h"
@@ -28,13 +29,11 @@ limitations under the License.
 #include "absl/types/span.h"
 #include "xla/hlo/analysis/alias_info.h"
 #include "xla/hlo/analysis/hlo_dataflow_analysis.h"
-#include "xla/hlo/analysis/hlo_ordering.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_module.h"
 #include "xla/service/hlo_buffer.h"
 #include "xla/service/hlo_value.h"
 #include "xla/shape_util.h"
-#include "xla/types.h"
 #include "xla/xla_data.pb.h"
 
 namespace xla {
@@ -116,9 +115,22 @@ class HloAliasAnalysis {
   // A map indicating which buffer a value is contained in.
   absl::flat_hash_map<const HloValue*, HloBuffer*> value_to_buffer_;
 
+  // Returns whether the computation has a chain of callers where every
+  // computation up to the entry computation has exactly one caller.
+  bool HasLinearCallerChainToRoot(const HloComputation* computation) const;
+
+  // Returns whether the set of given buffers are allowed to exist at one
+  // position before inlining.
+  bool MultipleBuffersAllowedBeforeInlining(
+      absl::Span<const HloBuffer* const> buffers) const;
+
   // A lazily constructed vector containing all HloBuffers sorted by
   // HloBuffer::Id.
   std::vector<HloBuffer> buffers_;
+
+  // Cache for whether an HloComputation has a single-caller chain to root.
+  mutable absl::flat_hash_map<const HloComputation*, bool>
+      has_linear_caller_chain_to_root_;
 };
 
 }  // namespace xla
diff --git a/third_party/xla/xla/hlo/analysis/hlo_alias_analysis_test.cc b/third_party/xla/xla/hlo/analysis/hlo_alias_analysis_test.cc
index 1f87c30b7ba7d0..4934f47a7fc879 100644
--- a/third_party/xla/xla/hlo/analysis/hlo_alias_analysis_test.cc
+++ b/third_party/xla/xla/hlo/analysis/hlo_alias_analysis_test.cc
@@ -1418,5 +1418,66 @@ ENTRY main {
             analysis.GetUniqueBufferAt(call, {1}));
 }
 
+TEST_F(HloAliasAnalysisTest,
+       AsyncComputationMultipleCallersMultipleBuffersAllowed) {
+  absl::string_view hlo_string = R"(
+HloModule Module
+
+async_computation {
+  ROOT p = f32[16] parameter(0)
+}
+
+ENTRY main {
+  p0 = f32[16] parameter(0)
+  p1 = f32[16] parameter(1)
+  async-start.0 = ((f32[16]), f32[16], s32[]) async-start(p0), calls=async_computation
+  async-done.0 = f32[16] async-done(async-start.0), calls=async_computation
+  async-start.1 = ((f32[16]), f32[16], s32[]) async-start(p1), calls=async_computation
+  async-done.1 = f32[16] async-done(async-start.1), calls=async_computation
+  ROOT tuple = (f32[16], f32[16]) tuple(async-done.0, async-done.1)
+}
+)";
+  ASSERT_OK_AND_ASSIGN(module_, ParseAndReturnVerifiedModule(hlo_string));
+  HloAliasAnalysis& analysis = RunAnalysis();
+
+  HloComputation* async_computation =
+      module_->GetComputationWithName("async_computation");
+  ASSERT_NE(async_computation, nullptr);
+  const HloInstruction* param = async_computation->GetInstructionWithName("p");
+  ASSERT_NE(param, nullptr);
+
+  const HloInstruction* p0 =
+      module_->entry_computation()->GetInstructionWithName("p0");
+  const HloInstruction* p1 =
+      module_->entry_computation()->GetInstructionWithName("p1");
+  ASSERT_NE(p0, nullptr);
+  ASSERT_NE(p1, nullptr);
+
+  const HloBuffer& buffer0 = analysis.GetUniqueBufferAt(p0);
+  const HloBuffer& buffer1 = analysis.GetUniqueBufferAt(p1);
+  EXPECT_NE(&buffer0, &buffer1);
+
+  // Because async_computation has two callers operating on two different
+  // buffers (p0 and p1), multiple buffers are allowed at positions inside and
+  // flowing out of async_computation without triggering a crash.
+  std::vector<const HloBuffer*> param_buffers =
+      analysis.ComputeBuffersAt(param);
+  EXPECT_THAT(param_buffers, UnorderedElementsAre(&buffer0, &buffer1));
+
+  const HloInstruction* async_done0 =
+      module_->entry_computation()->GetInstructionWithName("async-done.0");
+  ASSERT_NE(async_done0, nullptr);
+  std::vector<const HloBuffer*> done0_buffers =
+      analysis.ComputeBuffersAt(async_done0);
+  EXPECT_THAT(done0_buffers, UnorderedElementsAre(&buffer0, &buffer1));
+
+  const HloInstruction* tuple =
+      module_->entry_computation()->GetInstructionWithName("tuple");
+  ASSERT_NE(tuple, nullptr);
+  std::vector<const HloBuffer*> tuple_elem0_buffers =
+      analysis.ComputeBuffersAt(tuple, {0});
+  EXPECT_THAT(tuple_elem0_buffers, UnorderedElementsAre(&buffer0, &buffer1));
+}
+
 }  // namespace
 }  // namespace xla
diff --git a/third_party/xla/xla/hlo/analysis/while_loop_analysis.cc b/third_party/xla/xla/hlo/analysis/while_loop_analysis.cc
index e80c4c5e3ba9b6..cb06ad3a7858c9 100644
--- a/third_party/xla/xla/hlo/analysis/while_loop_analysis.cc
+++ b/third_party/xla/xla/hlo/analysis/while_loop_analysis.cc
@@ -774,7 +774,7 @@ optional<Range> MatchLoopRangeWithKnownValues(
   int64_t trip_count_step = 0;
   if (!Match(while_body_indvar_update,
              m::AddAnyOrder(m::Op().Is(while_body_indvar),
-                            m::Op(&trip_count_increase_step_instr)))) {
+                            m::Constant(&trip_count_increase_step_instr)))) {
     if (trip_count_increase_step_instr == nullptr) {
       VLOG(2) << "Pattern-match failed: induction variable is not getting "
                  "updated by an add operation: "
diff --git a/third_party/xla/xla/hlo/analysis/while_loop_analysis_test.cc b/third_party/xla/xla/hlo/analysis/while_loop_analysis_test.cc
index 972ca305a6cf61..c89c2df07c83b9 100644
--- a/third_party/xla/xla/hlo/analysis/while_loop_analysis_test.cc
+++ b/third_party/xla/xla/hlo/analysis/while_loop_analysis_test.cc
@@ -1066,7 +1066,7 @@ TEST_F(WhileLoopAnalysisTest, GetIndvarIndexShouldWorkWhenParamIsCopied) {
 }
 
 TEST_F(WhileLoopAnalysisTest,
-       MatchTrivialLoopCountFailsWhenIndvarIsNotIncrementedByConstant) {
+       MatchTrivialLoopFailsWhenIndvarIsNotIncrementedByConstant) {
   absl::string_view hlo_with_constant = R"(
   HloModule test
   body {
@@ -1125,6 +1125,7 @@ TEST_F(WhileLoopAnalysisTest,
       MatchTrivialLoopTripCount(while_op_without_constant, 0,
                                 LiteralUtil::CreateR0<int32_t>(0));
   EXPECT_EQ(trip_count_without_constant, std::nullopt);
+  EXPECT_EQ(MatchTrivialLoopRange(while_op_without_constant), std::nullopt);
 }
 
 TEST_F(WhileLoopAnalysisTest,
diff --git a/third_party/xla/xla/mlir_hlo/deallocation/transforms/passes.h b/third_party/xla/xla/mlir_hlo/deallocation/transforms/passes.h
index ebfc7e5c588578..c44a7c370dd70f 100644
--- a/third_party/xla/xla/mlir_hlo/deallocation/transforms/passes.h
+++ b/third_party/xla/xla/mlir_hlo/deallocation/transforms/passes.h
@@ -29,11 +29,6 @@ namespace deallocation {
 #define GEN_PASS_DECL
 #include "deallocation/transforms/passes.h.inc"
 
-// TODO(b/397167511): Remove legacy wrapper once callers are migrated.
-inline std::unique_ptr<mlir::Pass> createBufferDeallocationPass() {
-  return createBufferDeallocation();
-}
-
 #define GEN_PASS_REGISTRATION
 #include "deallocation/transforms/passes.h.inc"
 
diff --git a/third_party/xla/xla/mlir_hlo/mhlo/transforms/passes.h b/third_party/xla/xla/mlir_hlo/mhlo/transforms/passes.h
index 1f235c3e4fb44a..b45c9d6d7a47fb 100644
--- a/third_party/xla/xla/mlir_hlo/mhlo/transforms/passes.h
+++ b/third_party/xla/xla/mlir_hlo/mhlo/transforms/passes.h
@@ -44,12 +44,6 @@ ChloLegalizeToHighLevelMhloPassOptions getDefaultChloToHighLevelMhloOptions();
 /// Returns options for the ChloLegalizeToHighLevelMhloPass for the GPU backend.
 ChloLegalizeToHighLevelMhloPassOptions getGpuChloToHighLevelMhloOptions();
 
-// TODO(b/397167511): Remove legacy wrapper once callers are migrated.
-inline std::unique_ptr<mlir::Pass>
-createLegalizeTrigonometricToApproximationPass() {
-  return createLegalizeTanhToApproximationPass();
-}
-
 // TODO(b/397167511): Remove legacy wrapper once callers are migrated.
 inline std::unique_ptr<mlir::Pass> createExpandHloTuplesPass(
     const std::string& entryFunctionName) {
diff --git a/third_party/xla/xla/mlir_hlo/transforms/gpu_passes.h b/third_party/xla/xla/mlir_hlo/transforms/gpu_passes.h
index 2711068d5b2c97..616f852070ed7e 100644
--- a/third_party/xla/xla/mlir_hlo/transforms/gpu_passes.h
+++ b/third_party/xla/xla/mlir_hlo/transforms/gpu_passes.h
@@ -37,22 +37,6 @@ struct Chipset;
 // 'gpu.launc_func' ops during the fusion rewrite pass above.
 ArrayAttr getWrittenOperandsAttribute(Operation* op);
 
-/// Pass that transforms gpu modules in standard dialect to NVVM.
-inline std::unique_ptr<mlir::Pass> createGpuKernelToNvvmPass(
-    bool useBarePtrCallConv = false) {
-  GpuKernelToNVVMPassOptions options;
-  options.useBarePtrCallConv = useBarePtrCallConv;
-  return createGpuKernelToNVVMPass(options);
-}
-
-/// Pass that transforms gpu modules in standard dialect to ROCDL.
-inline std::unique_ptr<mlir::Pass> createGpuKernelToRocdlPass(
-    const std::string& chipset = "gfx000") {
-  GpuKernelToROCDLPassOptions options;
-  options.chipset = chipset;
-  return createGpuKernelToROCDLPass(options);
-}
-
 #define GEN_PASS_REGISTRATION
 #include "transforms/gpu_passes.h.inc"
 
diff --git a/third_party/xla/xla/mlir_hlo/transforms/passes.h b/third_party/xla/xla/mlir_hlo/transforms/passes.h
index 7cf86ceb6c6026..ce5a0740baf80b 100644
--- a/third_party/xla/xla/mlir_hlo/transforms/passes.h
+++ b/third_party/xla/xla/mlir_hlo/transforms/passes.h
@@ -64,14 +64,6 @@ inline std::unique_ptr<Pass> createTileLoopsPass(
 }
 
 namespace hlo {
-using mlir::createAllocToArgPass;
-using mlir::createGenericHostToLLVMPass;
-using mlir::createUnbufferizePass;
-
-inline std::unique_ptr<Pass> createOneShotBufferizePass() {
-  return mlir::createOneShotBufferize();
-}
-
 #define GEN_PASS_REGISTRATION
 #include "transforms/passes.h.inc"
 
diff --git a/third_party/xla/xla/service/elemental_ir_emitter.cc b/third_party/xla/xla/service/elemental_ir_emitter.cc
index 6f9fee1d215a3a..e49674dde9c412 100644
--- a/third_party/xla/xla/service/elemental_ir_emitter.cc
+++ b/third_party/xla/xla/service/elemental_ir_emitter.cc
@@ -2440,31 +2440,44 @@ llvm::Value* ElementalIrEmitter::EmitIntegerRemainder(llvm::Value* lhs,
       Select(has_int_min_overflow, GetZero(lhs->getType()), safe_rem));
 }
 
-llvm::Value* ElementalIrEmitter::EmitIntegerPow(llvm::Value* base,
-                                                llvm::Value* exponent,
+llvm::Value* ElementalIrEmitter::EmitIntegerPow(llvm::Value* lhs,
+                                                llvm::Value* rhs,
                                                 bool is_signed) {
   // Exponentiation by squaring:
   // https://en.wikipedia.org/wiki/Exponentiation_by_squaring;
   int bits = 6;  // Everything else would overflow for any exponent > 1, as 2^64
                  // is the larget possible exponent for a 64-bit integer, and
                  // that's 1 << 6.
-  llvm::Value* accumulator = llvm::ConstantInt::get(base->getType(), 1);
-  llvm::Value* one = llvm::ConstantInt::get(exponent->getType(), 1);
-  llvm::Value* zero = llvm::ConstantInt::get(exponent->getType(), 0);
+  llvm::Value* base = lhs;
+  llvm::Value* exponent = rhs;
+  llvm::Value* exp_one = llvm::ConstantInt::get(exponent->getType(), 1);
+  llvm::Value* exp_zero = llvm::ConstantInt::get(exponent->getType(), 0);
+  llvm::Value* base_one = llvm::ConstantInt::get(base->getType(), 1);
+  llvm::Value* base_zero = llvm::ConstantInt::get(base->getType(), 0);
+  llvm::Value* base_neg_one = llvm::ConstantInt::get(base->getType(), -1, true);
+  llvm::Value* accumulator = base_one;
   llvm::Value* original_base = base;
   llvm::Value* original_exponent = exponent;
 
   // Unroll the loop at compile time.
   for (int i = 0; i < bits; i++) {
-    accumulator =
-        b_->CreateSelect(b_->CreateICmpEQ(b_->CreateAnd(exponent, one), one),
-                         b_->CreateMul(accumulator, base), accumulator);
+    accumulator = b_->CreateSelect(
+        b_->CreateICmpEQ(b_->CreateAnd(exponent, exp_one), exp_one),
+        b_->CreateMul(accumulator, base), accumulator);
     base = b_->CreateMul(base, base);
     exponent = b_->CreateLShr(exponent, 1);
   }
+
+  llvm::Value* neg_one_base_result = b_->CreateSelect(
+      b_->CreateICmpEQ(b_->CreateAnd(original_exponent, exp_one), exp_one),
+      base_neg_one, base_one);
+  llvm::Value* neg_exp_res =
+      b_->CreateSelect(b_->CreateICmpEQ(original_base, base_neg_one),
+                       neg_one_base_result, base_zero);
   return b_->CreateSelect(
-      b_->CreateICmpSGE(original_exponent, zero), accumulator,
-      b_->CreateSelect(b_->CreateICmpEQ(original_base, one), one, zero));
+      b_->CreateICmpSGE(original_exponent, exp_zero), accumulator,
+      b_->CreateSelect(b_->CreateICmpEQ(original_base, base_one), base_one,
+                       neg_exp_res));
 }
 
 llvm::Value* ElementalIrEmitter::EmitIntegerMulhi(llvm::Value* lhs,
diff --git a/third_party/xla/xla/service/elemental_ir_emitter_test.cc b/third_party/xla/xla/service/elemental_ir_emitter_test.cc
index 70f0197aaae2f8..2cb64aa48fd0ac 100644
--- a/third_party/xla/xla/service/elemental_ir_emitter_test.cc
+++ b/third_party/xla/xla/service/elemental_ir_emitter_test.cc
@@ -620,5 +620,22 @@ ENTRY e {
                                                 /*arel=*/1e-3}));
 }
 
+TEST_F(ElementalIrEmitterExecutionTest, IntegerPowNegativeExponent) {
+  const std::string hlo_text = R"(
+HloModule IntegerPowNegativeExponent
+
+ENTRY main {
+  base = s32[8]{0} parameter(0)
+  exponent = s32[8]{0} parameter(1)
+  ROOT power = s32[8]{0} power(base, exponent)
+}
+)";
+
+  Literal base = LiteralUtil::CreateR1<int32_t>({-1, -1, -1, -1, 1, 2, -2, 3});
+  Literal exponent =
+      LiteralUtil::CreateR1<int32_t>({-1, -2, -3, -100, -5, -1, -1, -2});
+  RunTest(hlo_text, {&base, &exponent});
+}
+
 }  // namespace
 }  // namespace xla
diff --git a/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass.cc b/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass.cc
index f43d22cd851b5b..02337aaa512f05 100644
--- a/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass.cc
+++ b/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass.cc
@@ -297,7 +297,6 @@ ConfigAssigner::Options GetConfigAssignerOptions(
 CodegenOrchestrator::Options GetCodegenOrchestratorOptions(
     const DebugOptions& debug_options) {
   CodegenOrchestrator::Options options;
-  options.exclude_cublas_config = !debug_options.xla_gpu_cublas_fallback();
   if (!debug_options.xla_gpu_fail_ptx_compilation_on_register_spilling()) {
     options.allow_reg_spills_fn = [](const HloInstruction& instr,
                                      autotuner::Backend backend) {
@@ -335,6 +334,12 @@ ProfileOptions GetProfileOptions(const DebugOptions& debug_options,
 Autotuner::Options GetAutotunerOptions(const DebugOptions& debug_options,
                                        bool is_buffer_check_supported) {
   Autotuner::Options autotuner_options;
+  if (!debug_options.xla_gpu_cublas_fallback()) {
+    autotuner_options.excluded_backends.push_back(
+        autotuner::Backend::CUBLASLT_FISSION);
+    autotuner_options.excluded_backends.push_back(
+        autotuner::Backend::HIPBLASLT_FISSION);
+  }
   autotuner_options.correctness_check_options.enable_correctness_check =
       is_buffer_check_supported && debug_options.xla_gpu_autotune_level() >= 4;
   autotuner_options.correctness_check_options.relative_tolerance =
diff --git a/third_party/xla/xla/xla.proto b/third_party/xla/xla/xla.proto
index 5d8d11eb9987c1..b9d9e0374324f3 100644
--- a/third_party/xla/xla/xla.proto
+++ b/third_party/xla/xla/xla.proto
@@ -1827,7 +1827,12 @@ message DebugOptions {
   // Gives more compilation test coverage but might increase compilation time.
   optional bool xla_compile_all_supported_configs = 539;
 
-  // Next id: 541
+  // Internal debug/testing flag to enable cross-host one-shot kernel for
+  // collective operations.
+  repeated CollectiveOpType xla_gpu_unsupported_use_cross_host_one_shot_kernel =
+      541;
+
+  // Next id: 542
 
   // Extra options to pass to the compilation backend (e.g. LLVM); specific
   // interpretation of these values is left to the backend.