PR tensorflow#43503: [ROCM] Important fixes for GpuBlasLt MatmulPlan cache

Imported from GitHub PR openxla/xla#43503

## Summary

This change prevents GPU BLASLt autotuning from blowing up the MatmulPlan cache by removing the selected algorithm from the plan cache key. MatmulPlans are now cached by the canonical GEMM key with the selected algorithm unset, while each plan caches its own algorithm list and currently selected algorithm.

The original problem came with a new GEMM autotuner which, instead of calling gpublas_lt directly for each supported algorithm, profiles tiny GpuExecutables. That is, during the autotuning, we go through CublasLtMatmulThunk::ExecuteOnStream(..) which then inserts new entries to MatmulPlan cache (since canonical HLOs differ by **selected algorithm**):
```
I0000 00:00:1779469188.517485 1847675 gpublas_lt_matmul_thunk.cc:140] 0x7f4d0403d9c0: Adding new grouped MatmulPlan for stream: 0x21d619e0 instr: (f16[2400,8]{1,0}, s8[117600]{0}) custom-call(f16[2400,9]{1,0}, f16[600,9,8]{2,1,0}, s32[600]{0}), custom_call_target="__cublas$lt$groupedMatmul", backend_config={"operation_queue_id":"0","force_earliest_schedule":false,"reification_cost":[],"device_type":"DEVICE_TYPE_INVALID","grouped_gemm_backend_config":{"gemm_backend_config":{"selected_algorithm":"0","alpha_real":1,"beta":0,"dot_dimension_numbers":{"lhs_contracting_dimensions":["1"],"rhs_contracting_dimensions":["1"],"lhs_batch_dimensions":[],"rhs_batch_dimensions":[]},"alpha_imag":0,"epilogue":"DEFAULT","grad_x":false,"grad_y":false,"damax_output":false,"autotune_workspace_size":"117600","scale_mode":0},"ragged_dot_dimension_numbers":{"dot_dimension_numbers":{"lhs_contracting_dimensions":["1"],"rhs_contracting_dimensions":["1"],"lhs_batch_dimensions":[],"rhs_batch_dimensions":[]},"lhs_ragged_dimensions":["0"],"rhs_group_dimensions":["0"]}}}
I0000 00:00:1779469188.957549 1847675 gpu_blas_lt.cc:314] 0x7f4f68001478: Plan created: cache size: 1
I0000 00:00:1779469190.125376 1847675 gpublas_lt_matmul_thunk.cc:140] 0x7f4d0803d9d0: Adding new grouped MatmulPlan for stream: 0x21d619e0 instr: (f16[2400,8]{1,0}, s8[117600]{0}) custom-call(f16[2400,9]{1,0}, f16[600,9,8]{2,1,0}, s32[600]{0}), custom_call_target="__cublas$lt$groupedMatmul", backend_config={"operation_queue_id":"0","force_earliest_schedule":false,"reification_cost":[],"device_type":"DEVICE_TYPE_INVALID","grouped_gemm_backend_config":{"gemm_backend_config":{"selected_algorithm":"1","alpha_real":1,"beta":0,"dot_dimension_numbers":{"lhs_contracting_dimensions":["1"],"rhs_contracting_dimensions":["1"],"lhs_batch_dimensions":[],"rhs_batch_dimensions":[]},"alpha_imag":0,"epilogue":"DEFAULT","grad_x":false,"grad_y":false,"damax_output":false,"autotune_workspace_size":"117600","scale_mode":0},"ragged_dot_dimension_numbers":{"dot_dimension_numbers":{"lhs_contracting_dimensions":["1"],"rhs_contracting_dimensions":["1"],"lhs_batch_dimensions":[],"rhs_batch_dimensions":[]},"lhs_ragged_dimensions":["0"],"rhs_group_dimensions":["0"]}}}
I0000 00:00:1779469202.151521 1847675 gpu_blas_lt.cc:314] 0x7f4f68001478: Plan created: cache size: 2
```
This might not be that bad for regular gemms, but can slow down the autotuning grouped gemm significantly. Furthermore, the MatmulPlan cache keeps a list of **128 MatmulPlans** for each gemm config after autotuning.

## Details

- Introduces two-level caching for GPU BLASLt matmul execution: the outer cache reuses MatmulPlans across autotuned algorithm choices, and the inner cache keeps the retrieved algorithms and selected algorithm on the plan.
- Updates thunk emission and execution to pass the selected algorithm separately from the canonical plan key.
- Applies the cache flow across CUDA and ROCm BLASLt implementations, including grouped GEMM handling.
- Adds coverage asserting that autotuned GEMM and grouped GEMM create a single MatmulPlan cache entry instead of one entry per selected algorithm.

## Test Plan

- Added/updated `GpuBlasLtMatmulThunkTest` coverage for shared MatmulPlans and autotuned cache behavior.

Copybara import of the project:

--
9a670c13c02ca01650590a0ac905f9991b4d7b5c by Pavel Emeliyanenko <pavel.emeliyanenko@amd.com>:

redesigning blaslt cache

Revert "redesigning blaslt cache"

This reverts commit 5a96365fc83024a011c5a829926aed5cbd0d8594.

moving towards cached algorithms inside matmul plan

adding setcachedalgorithm

added cached index

improving set cached algorithm

simplified set algorithm

fixing cuda build

cleanup

fixing macros

Revert "redesigning blaslt cache"

This reverts commit 8a7481576e5fff2b13d39b646d3d607db8bb9628.

fixing thunk emitter cache key

just skip the test to check if it hangs

reenabled subtests

disabled all

enabled 2 subtests

disabled GemmWithAutotuneOneCacheEntry

disabled functional test

enabled autotune tests

enabled only autotuneonecache

enabled just single test

running single test

fix

debugging why matmul_test hangs on CI

reenabled again canonicalhlo string

use workspace_size in the if statement

restored ci_build

--
ad8dfba320b4dc4de038570a9edc7d51a4746099 by Pavel Emeliyanenko <pavel.emeliyanenko@amd.com>:

reenabled grouped gemm test

--
685309d9109ba3dfcea95976e2bb3b9b144f05be by Pavel Emeliyanenko <pavel.emeliyanenko@amd.com>:

reduced the number of parallel threads

Merging this change closes tensorflow#43503

PiperOrigin-RevId: 932409289

Author: pemeliya

third_party/xla/xla/backends/gpu/runtime/BUILD

@@ -1475,6 +1475,7 @@ xla_test(
         "//xla/runtime:buffer_use",
         "//xla/service:buffer_assignment",
         "//xla/service:executable",
+        "//xla/service:hlo_module_config",
         "//xla/service:platform_util",
         "//xla/service:shaped_slice",
         "//xla/service/gpu:backend_configs_cc",

third_party/xla/xla/backends/gpu/runtime/gpublas_lt_matmul_thunk.cc

@@ -129,6 +129,7 @@ absl::Status CublasLtMatmulThunk::ExecuteOnStream(const ExecuteParams& params) {
   }
   return plan->ExecuteOnStream(params.stream, args, /*profile_result*/ nullptr);
 }
+
 absl::StatusOr<se::gpu::BlasLt::MatmulPlan*>
 CublasLtMatmulThunk::GetCachedMatmulPlan(const ExecuteParams& params) {
   ASSIGN_OR_RETURN(auto* blas_lt,
@@ -139,32 +140,19 @@ CublasLtMatmulThunk::GetCachedMatmulPlan(const ExecuteParams& params) {
             << " instr: " << canonical_hlo_;
 
     ASSIGN_OR_RETURN(auto plan, std::visit(
-                                    [&](const auto& gemm_config) {
+                                    [&](auto&& gemm_config) {
                                       return blas_lt->GetMatmulPlan(gemm_config,
                                                                     epilogue_);
                                     },
                                     gemm_config_));
-
-    // Set the workspace size to the size that was used for autotuning, so
-    // algorithm index will be the same as returned by GetAlgorithms called
-    // during autotuning.
-    int64_t max_workspace = autotune_workspace_size_;
-
-    // If autotuning is disabled, there is no point on retrieving all
-    // algorithms, it's enough to get the default one only.
-    int64_t num_algorithms =
-        algorithm_idx_ == 0 ? 1 : GemmConfig::kNumAlgorithms;
-    ASSIGN_OR_RETURN(auto algorithms,
-                     plan->GetAlgorithms(num_algorithms, max_workspace));
-
-    if (algorithms.empty()) {
-      return absl::InternalError(
-          "Failed to get a MatmulPlan: no valid algorithm found.");
-    }
-    RETURN_IF_ERROR(plan->SetAlgorithm(algorithms[algorithm_idx_]));
     return std::move(plan);
   };
-  return blas_lt->GetOrCreateMatmulPlan(canonical_hlo_, create);
+  // If autotuning is disabled, there is no point on retrieving all
+  // algorithms, it's enough to get the default one only.
+  size_t num_algorithms = algorithm_idx_ == 0 ? 1 : GemmConfig::kNumAlgorithms;
+  return blas_lt->GetOrCreateMatmulPlanWithAlgorithm(
+      canonical_hlo_, create, algorithm_idx_, num_algorithms,
+      autotune_workspace_size_);
 }
 
 absl::Status CublasLtMatmulThunk::Initialize(const InitializeParams& params) {
@@ -216,7 +204,7 @@ absl::StatusOr<ThunkProto> CublasLtMatmulThunk::ToProto() const {
       proto.mutable_cublas_lt_matmul_thunk();
 
   RETURN_IF_ERROR(std::visit(
-      [&](const auto& gemm_config) {
+      [&](auto&& gemm_config) {
         using T = std::decay_t<decltype(gemm_config)>;
         if constexpr (std::is_same_v<T, se::gpu::GroupedGemmConfig>) {
           *cublas_lt_matmul_thunk->mutable_grouped_gemm_config() =

third_party/xla/xla/backends/gpu/runtime/gpublas_lt_matmul_thunk_test.cc

@@ -45,13 +45,15 @@ limitations under the License.
 #include "xla/error_spec.h"
 #include "xla/executable_run_options.h"
 #include "xla/hlo/ir/hlo_instruction.h"
+#include "xla/hlo/ir/hlo_module.h"
 #include "xla/hlo/ir/hlo_print_options.h"
 #include "xla/runtime/buffer_use.h"
 #include "xla/service/buffer_assignment.h"
 #include "xla/service/gpu/backend_configs.pb.h"
 #include "xla/service/gpu/buffer_allocations.h"
 #include "xla/service/gpu/cublas_cudnn.h"
 #include "xla/service/gpu/matmul_utils.h"
+#include "xla/service/hlo_module_config.h"
 #include "xla/service/platform_util.h"
 #include "xla/service/service_executable_run_options.h"
 #include "xla/service/shaped_slice.h"
@@ -199,7 +201,6 @@ void GpuBlasLtMatmulThunkTest::CreateExecuteThunksFromHLO(
     se::StreamExecutor* executor, absl::string_view hlo_string) {
   TF_ASSERT_OK_AND_ASSIGN(auto module,
                           this->ParseAndReturnVerifiedModule(hlo_string));
-
   GemmRewriterOptions options;
   TF_ASSERT_OK_AND_ASSIGN(
       bool changed,
@@ -211,7 +212,6 @@ void GpuBlasLtMatmulThunkTest::CreateExecuteThunksFromHLO(
 
   GpuBlasLtThunkBuilder builder(executor, gpu_comp(executor));
   std::vector<std::unique_ptr<CublasLtMatmulThunk>> gemm_thunks;
-
   for (auto* instr : module->entry_computation()->instructions()) {
     if (IsCublasLtMatmul(*instr)) {
       TF_ASSERT_OK_AND_ASSIGN(auto thunk, builder.CreateThunk(instr));
@@ -329,6 +329,71 @@ TEST_F(GpuBlasLtMatmulThunkTest, SharedMatmulPlansFunctional) {
   EXPECT_EQ(blas_lt->GetMatmulPlanCacheSize(), 2);
 }
 
+TEST_F(GpuBlasLtMatmulThunkTest, GemmWithAutotuneOneCacheEntry) {
+  auto* exec = default_exec();
+  auto* blas_lt = exec->AsBlas()->GetBlasLt();
+  EXPECT_NE(blas_lt, nullptr);
+  blas_lt->ClearMatmulPlanCache();
+
+  constexpr absl::string_view simple_gemm_hlo = R"(
+HloModule SimpleGemm
+
+ENTRY AddDotsFunc {
+  x = f32[128,128] parameter(0)
+  y = f32[128,128] parameter(1)
+  ROOT dot_a = f32[128,128] dot(x, y), lhs_contracting_dims={1}, rhs_contracting_dims={0}
+})";
+
+  HloModuleConfig config = GetModuleConfigForTest();
+  auto& debug_opts = config.mutable_debug_options();
+  debug_opts.set_xla_gpu_enable_cublaslt(true);
+  debug_opts.set_xla_gpu_autotune_level(1);
+  debug_opts.set_xla_gpu_enable_triton_gemm(false);
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<HloModule> module,
+      ParseAndReturnVerifiedModule(simple_gemm_hlo, config));
+  EXPECT_TRUE(RunAndCompare(std::move(module), ErrorSpec{1e-3, 1e-3}));
+  EXPECT_EQ(blas_lt->GetMatmulPlanCacheSize(), 1);
+}
+
+TEST_F(GpuBlasLtMatmulThunkTest, GroupedGemmWithAutotuneOneCacheEntry) {
+  auto* rocm = gpu_comp().rocm_compute_capability();
+  if (rocm == nullptr || !rocm->gfx9_mi300_series()) {
+    GTEST_SKIP() << "Grouped GEMM is only supported on ROCm gfx942 or gfx950";
+  }
+
+  auto* exec = default_exec();
+  auto* blas_lt = exec->AsBlas()->GetBlasLt();
+  EXPECT_NE(blas_lt, nullptr);
+  blas_lt->ClearMatmulPlanCache();
+
+  constexpr absl::string_view grouped_gemm_hlo = R"(
+HloModule GroupedGemm
+
+ENTRY AddRaggedDotsFunc {
+  p0 = f16[64,9]{1,0} parameter(0)
+  p1 = f16[2,9,8]{2,1,0} parameter(1)
+  p2 = s32[2] constant({16, 48})
+  ROOT ragged-dot = f16[64,8]{1,0} ragged-dot(p0, p1, p2),
+                    lhs_contracting_dims={1}, rhs_contracting_dims={1},
+                    lhs_ragged_dims={0}, rhs_group_dims={0}
+})";
+
+  HloModuleConfig config = GetModuleConfigForTest();
+  auto& debug_opts = config.mutable_debug_options();
+  debug_opts.set_xla_gpu_autotune_level(1);
+  debug_opts.set_xla_gpu_enable_cublaslt(true);
+
+  debug_opts.set_xla_gpu_enable_triton_gemm(false);
+  debug_opts.set_xla_gpu_experimental_use_ragged_dot_grouped_gemm(true);
+  TF_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<HloModule> grouped_module,
+      ParseAndReturnVerifiedModule(grouped_gemm_hlo, config));
+  EXPECT_TRUE(RunAndCompare(std::move(grouped_module), ErrorSpec{1e-4, 1e-5}));
+  EXPECT_EQ(blas_lt->GetMatmulPlanCacheSize(), 1);
+}
+
 // Mock BlasLt interface to test only the cache function
 struct MockBlasLt : public se::gpu::BlasLt {
   absl::Status Init() override { return absl::OkStatus(); }
@@ -346,19 +411,40 @@ struct MockBlasLt : public se::gpu::BlasLt {
   ~MockBlasLt() override = default;
 };
 
+struct MockMatmulPlan : public se::gpu::BlasLt::MatmulPlan {
+  absl::StatusOr<std::vector<se::gpu::BlasLt::MatmulAlgorithm>> GetAlgorithms(
+      size_t max_algorithm_count, size_t max_workspace_size) const override {
+    se::gpu::BlasLt::MatmulAlgorithm algo;
+    return std::vector<se::gpu::BlasLt::MatmulAlgorithm>{algo};
+  }
+
+  absl::Status SetAlgorithm(const se::gpu::BlasLt::MatmulAlgorithm&) override {
+    return absl::OkStatus();
+  }
+
+  absl::Status ExecuteOnStream(
+      se::Stream* stream, const se::gpu::BlasLt::MemoryArgs& args,
+      se::blas::ProfileResult* profile_result) const override {
+    return absl::OkStatus();
+  }
+  ~MockMatmulPlan() override = default;
+};
+
 TEST_F(GpuBlasLtMatmulThunkTest, CacheUnitTest) {
   auto thread_func = [&](MockBlasLt* blas_lt, const std::string& key,
                          int sleep_ms) -> absl::Status {
     auto create_func = [&]() -> absl::StatusOr<se::gpu::BlasLt::MatmulPlanPtr> {
       // We don't care about creation of matmul plans -> emulate it with a sleep
       absl::SleepFor(absl::Milliseconds(sleep_ms));
-      return se::gpu::BlasLt::MatmulPlanPtr{};
+      return std::make_unique<MockMatmulPlan>();
     };
 
-    return blas_lt->GetOrCreateMatmulPlan(key, create_func).status();
+    return blas_lt
+        ->GetOrCreateMatmulPlanWithAlgorithm(key, create_func, 0, 1, 0)
+        .status();
   };  // thread_func
 
-  const int num_blas_lts = 30, num_streams = 30,
+  const int num_blas_lts = 8, num_streams = 8,
             total = num_blas_lts * num_streams, mod = 11;
 
   std::vector<absl::Status> results(total);
@@ -573,9 +659,13 @@ static se::StreamExecutor* GpuExecutor() {
 }
 
 // Returns true if the GPU supports CUDA graph tracing (requires CUDA 12.3+).
-static bool SupportsCudaGraphTracing(const se::StreamExecutor* executor) {
+static bool SupportsGpuGraphTracing(const se::StreamExecutor* executor) {
   const auto& desc = executor->GetDeviceDescription();
-  const auto* cuda_cc = desc.gpu_compute_capability().cuda_compute_capability();
+  const auto& gpu_cc = desc.gpu_compute_capability();
+  if (gpu_cc.IsRocm()) {
+    return true;
+  }
+  const auto* cuda_cc = gpu_cc.cuda_compute_capability();
   if (cuda_cc == nullptr) {
     return false;
   }
@@ -598,8 +688,8 @@ class CublasLtMatmulThunkCmdBufTest : public ::testing::Test {
 
   void SetUp() override {
     executor_ = GpuExecutor();
-    if (!SupportsCudaGraphTracing(executor_)) {
-      GTEST_SKIP() << "CUDA graph tracing is not supported";
+    if (!SupportsGpuGraphTracing(executor_)) {
+      GTEST_SKIP() << "GPU graph tracing is not supported";
     }
 
     TF_ASSERT_OK_AND_ASSIGN(stream_, executor_->CreateStream());

third_party/xla/xla/service/gpu/backend_configs.proto

@@ -118,11 +118,10 @@ message GemmBackendConfig {
   int32 scale_mode = 21;
 }
 
-// Backend config for the GEMM operation running through cuBLASLt.
+// Backend config for the grouped (ragged-dot) GEMM operation running through
+// cuBLASLt.
 message GroupedGemmBackendConfig {
-  oneof gemm_config {
-    GemmBackendConfig gemm_backend_config = 1;
-  }
+  GemmBackendConfig gemm_backend_config = 1;
 
   xla.RaggedDotDimensionNumbers ragged_dot_dimension_numbers = 2;
 }

third_party/xla/xla/service/gpu/thunk_emitter.cc

@@ -321,6 +321,22 @@ ThunkSequence FlattenThunkSequence(std::vector<ThunkSequence>&& sequences) {
   return result;
 }
 
+absl::StatusOr<std::string> CanonicalGemmHlo(
+    const HloCustomCallInstruction* instr) {
+  ASSIGN_OR_RETURN(auto gpu_config, instr->backend_config<GpuBackendConfig>());
+
+  auto* gemm_config = gpu_config.has_grouped_gemm_backend_config()
+                          ? gpu_config.mutable_grouped_gemm_backend_config()
+                                ->mutable_gemm_backend_config()
+                          : gpu_config.mutable_gemm_backend_config();
+
+  // Clear algorithm-specific fields from the cache key
+  gemm_config->clear_selected_algorithm();
+  gemm_config->set_autotune_workspace_size(0);
+  return instr->ToString(HloPrintOptions::Fingerprint()) +
+         BackendConfigWrapper(gpu_config).GetRawString();
+}
+
 }  // namespace
 
 ThunkEmitter::ThunkEmitter(
@@ -568,8 +584,8 @@ absl::StatusOr<ThunkSequence> ThunkEmitter::EmitCublasLtMatmulThunk(
                    gpublas_lt::AsBlasLtEpilogue(epilogue));
   Thunk::ThunkInfo thunk_info = Thunk::ThunkInfo::WithProfileAnnotation(
       instr, ir_emitter_context_->GetNextThunkId());
-  std::string canonical_hlo = instr->ToString(
-      HloPrintOptions::Fingerprint().set_print_backend_config(true));
+
+  ASSIGN_OR_RETURN(std::string canonical_hlo, CanonicalGemmHlo(instr));
   auto thunk = std::make_unique<CublasLtMatmulThunk>(
       std::move(thunk_info), std::move(canonical_hlo), std::move(gemm_config),
       blas_lt_epilogue, algorithm, config.autotune_workspace_size(), a, b, c, d,
@@ -658,8 +674,7 @@ absl::StatusOr<ThunkSequence> ThunkEmitter::EmitCublasLtMatmulThunkF8(
                    gpublas_lt::AsBlasLtEpilogue(epilogue));
   Thunk::ThunkInfo thunk_info = Thunk::ThunkInfo::WithProfileAnnotation(
       instr, ir_emitter_context_->GetNextThunkId());
-  std::string canonical_hlo = instr->ToString(
-      HloPrintOptions::Fingerprint().set_print_backend_config(true));
+  ASSIGN_OR_RETURN(std::string canonical_hlo, CanonicalGemmHlo(instr));
   auto thunk = std::make_unique<CublasLtMatmulThunk>(
       std::move(thunk_info), std::move(canonical_hlo), std::move(gemm_config),
       blas_lt_epilogue, algorithm, config.autotune_workspace_size(), a, b, c, d,
@@ -734,8 +749,7 @@ absl::StatusOr<ThunkSequence> ThunkEmitter::EmitCublasLtGroupedMatmulThunk(
 
   Thunk::ThunkInfo thunk_info = Thunk::ThunkInfo::WithProfileAnnotation(
       instr, ir_emitter_context_->GetNextThunkId());
-  std::string canonical_hlo = instr->ToString(
-      HloPrintOptions::Fingerprint().set_print_backend_config(true));
+  ASSIGN_OR_RETURN(std::string canonical_hlo, CanonicalGemmHlo(instr));
 
   auto thunk = std::make_unique<CublasLtMatmulThunk>(
       std::move(thunk_info), std::move(canonical_hlo), std::move(gemm_config),
@@ -788,8 +802,7 @@ absl::StatusOr<ThunkSequence> ThunkEmitter::EmitCublasLtMatmulThunkMx(
                    gpublas_lt::AsBlasLtEpilogue(epilogue));
   Thunk::ThunkInfo thunk_info = Thunk::ThunkInfo::WithProfileAnnotation(
       instr, ir_emitter_context_->GetNextThunkId());
-  std::string canonical_hlo = instr->ToString(
-      HloPrintOptions::Fingerprint().set_print_backend_config(true));
+  ASSIGN_OR_RETURN(std::string canonical_hlo, CanonicalGemmHlo(instr));
   auto thunk = std::make_unique<CublasLtMatmulThunk>(
       std::move(thunk_info), std::move(canonical_hlo), std::move(gemm_config),
       blas_lt_epilogue, algorithm, config.autotune_workspace_size(), a, b, c, d,

third_party/xla/xla/stream_executor/cuda/cuda_blas_lt.cc

@@ -458,7 +458,7 @@ absl::Status BlasLt::MatmulPlan::ExecuteOnStream(
     blas::ProfileResult* profile_result) const {
   if (!algorithm_.has_value()) {
     return absl::InternalError(
-        "Algorithm must be set before calling DoMatMul!");
+        "Algorithm must be set before calling ExecuteOnStream!");
   }
   DeviceAddressBase a = args.a, b = args.b;
   DeviceAddressBase a_scale = args.a_scale, b_scale = args.b_scale;
@@ -474,21 +474,20 @@ absl::Status BlasLt::MatmulPlan::ExecuteOnStream(
   }
 
   void* workspace_addr = nullptr;
-  uint64_t workspace_size = algorithm_->workspace_size;
-  if (workspace_size > 0) {
+  uint64_t workspace_size = workspace_size_;
+  if (workspace_size_ > 0) {
     if (args.scratch_allocator != nullptr) {
       ASSIGN_OR_RETURN(DeviceAddress<uint8_t> alloc,
-                       args.scratch_allocator->AllocateBytes(workspace_size));
+                       args.scratch_allocator->AllocateBytes(workspace_size_));
       workspace_addr = gpu::GpuMemoryMutable(&alloc);
     } else {
       workspace_addr = args.workspace.opaque();
       size_t new_size = args.workspace.size();
-      TF_RET_CHECK(workspace_addr != nullptr && new_size >= workspace_size);
+      TF_RET_CHECK(workspace_addr != nullptr && new_size >= workspace_size_);
       workspace_size = new_size;
     }
   }
 
-  auto palgo = std::any_cast<cublasLtMatmulAlgo_t>(&algorithm_->opaque_algo);
   {
     absl::MutexLock lock(blas_lt_.mu_);
     TF_RET_CHECK(blas_lt_.handle_.get() != nullptr);
@@ -557,22 +556,19 @@ absl::Status BlasLt::MatmulPlan::ExecuteOnStream(
         blas_lt_.executor_->Activate();
 
     void* c_ptr = zero_beta_ ? nullptr : args.c.opaque();
-    if (palgo != nullptr) {
-      SE_CUBLAS_RETURN_IF_ERROR(cublasLtMatmul(
-          blas_lt_.handle_.get(), op_desc_.get(), &alpha_[0], a.opaque(),
-          a_desc_.get(), b.opaque(), b_desc_.get(), &beta_[0], c_ptr,
-          c_desc_.get(), args.d.opaque(), d_desc_.get(), palgo, workspace_addr,
-          workspace_size,
-          absl::bit_cast<CUstream>(stream->platform_specific_handle().stream)));
-    } else {
-      return absl::InternalError("cublaslt: Invalid algorithm type");
-    }
+    SE_CUBLAS_RETURN_IF_ERROR(cublasLtMatmul(
+        blas_lt_.handle_.get(), op_desc_.get(), &alpha_[0], a.opaque(),
+        a_desc_.get(), b.opaque(), b_desc_.get(), &beta_[0], c_ptr,
+        c_desc_.get(), args.d.opaque(), d_desc_.get(), &algorithm_.value(),
+        workspace_addr, workspace_size,
+        absl::bit_cast<CUstream>(stream->platform_specific_handle().stream)));
   }
 
   if (profile_result != nullptr) {
     ASSIGN_OR_RETURN(absl::Duration elapsed, timer->GetElapsedDuration());
     // set algorithm ID to be unique (otherwise it gets kDefaultAlgorithm ID)
-    profile_result->set_algorithm(reinterpret_cast<blas::AlgorithmType>(palgo));
+    profile_result->set_algorithm(
+        reinterpret_cast<blas::AlgorithmType>(&algorithm_.value()));
     profile_result->set_is_valid(true);
     profile_result->set_elapsed_time_in_ms(absl::ToDoubleMilliseconds(elapsed));
   }