Internal changes only.

runtime/executor/BUILD

@@ -245,8 +245,14 @@ cc_test(
 
 cc_library(
     name = "llm_litert_compiled_model_executor",
-    srcs = ["llm_litert_compiled_model_executor.cc"],
-    hdrs = ["llm_litert_compiled_model_executor.h"],
+    srcs = [
+        "llm_litert_compiled_model_executor.cc",
+        "metal_utils.cc",
+    ],
+    hdrs = [
+        "llm_litert_compiled_model_executor.h",
+        "metal_utils.h",
+    ],
     deps = [
         ":common_utils",
         ":executor_settings_base",
@@ -270,22 +276,8 @@ cc_library(
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:string_view",
         "@com_google_absl//absl/types:span",
-        "@litert//litert/c/internal:litert_external_litert_buffer_context",  # buildcleaner: keep
-        "@litert//litert/cc:litert_common",
-        "@litert//litert/cc:litert_compiled_model",
-        "@litert//litert/cc:litert_element_type",
-        "@litert//litert/cc:litert_environment",
-        "@litert//litert/cc:litert_expected",
-        "@litert//litert/cc:litert_layout",
-        "@litert//litert/cc:litert_macros",
-        "@litert//litert/cc:litert_model",
         "@litert//litert/cc:litert_model_types",
-        "@litert//litert/cc:litert_options",
-        "@litert//litert/cc:litert_ranked_tensor_type",
-        "@litert//litert/cc:litert_tensor_buffer",
         "@litert//litert/cc:litert_tensor_buffer_types",
-        "@litert//litert/cc/options:litert_cpu_options",
-        "@litert//litert/cc/options:litert_runtime_options",
         "//runtime/components:model_resources",
         "//runtime/components:model_resources_litert_lm",
         "//runtime/components:model_resources_task",
@@ -301,7 +293,28 @@ cc_library(
         "//runtime/util:tensor_buffer_util",
         "@litert//tflite/delegates/xnnpack:xnnpack_delegate",
         "@litert//tflite/types:half",
-    ],
+    ] + select({
+        "@litert//litert:litert_link_capi_so": [
+            "@litert//litert/cc:litert_api_with_dynamic_runtime",
+        ],
+        "//conditions:default": [
+            "@litert//litert/c/internal:litert_external_litert_buffer_context",  # buildcleaner: keep
+            "@litert//litert/cc:litert_common",
+            "@litert//litert/cc:litert_compiled_model",
+            "@litert//litert/cc:litert_element_type",
+            "@litert//litert/cc:litert_environment",
+            "@litert//litert/cc:litert_environment_options",
+            "@litert//litert/cc:litert_expected",
+            "@litert//litert/cc:litert_layout",
+            "@litert//litert/cc:litert_macros",
+            "@litert//litert/cc:litert_model",
+            "@litert//litert/cc:litert_options",
+            "@litert//litert/cc:litert_ranked_tensor_type",
+            "@litert//litert/cc:litert_tensor_buffer",
+            "@litert//litert/cc/options:litert_cpu_options",
+            "@litert//litert/cc/options:litert_runtime_options",
+        ],
+    }),
 )
 
 cc_test(

runtime/executor/llm_litert_compiled_model_executor.cc

@@ -14,6 +14,10 @@
 
 #include "runtime/executor/llm_litert_compiled_model_executor.h"
 
+#if defined(__APPLE__)
+#include "runtime/executor/metal_utils.h"
+#endif
+
 #include <algorithm>
 #include <atomic>
 #include <cstdint>
@@ -143,13 +147,26 @@ absl::Status CopyKvCacheBuffers(
     const absl::flat_hash_map<absl::string_view, TensorBuffer>&
         src_kv_cache_buffers,
     const absl::flat_hash_map<absl::string_view, TensorBuffer>&
-        dst_kv_cache_buffers) {
+        dst_kv_cache_buffers,
+    void* command_queue = nullptr) {
   for (const auto& [name, src_buffer] : src_kv_cache_buffers) {
     if (!dst_kv_cache_buffers.contains(name)) {
       return absl::FailedPreconditionError(
           absl::StrCat("KV cache buffer ", name, " not found."));
     }
     const auto& dst_buffer = dst_kv_cache_buffers.at(name);
+
+#if defined(__APPLE__)
+    LITERT_ASSIGN_OR_RETURN(
+        bool metal_copied,
+        TryCopyKvCacheMetal(src_buffer, const_cast<TensorBuffer&>(dst_buffer),
+                            src_index_to_copy_on_prefill, decode_batch_size,
+                            command_queue));
+    if (metal_copied) {
+      continue;
+    }
+#endif
+
     LITERT_ASSIGN_OR_RETURN(auto src_buffer_lock_and_addr,
                             TensorBufferScopedLock::Create(
                                 src_buffer, TensorBuffer::LockMode::kRead));
@@ -542,9 +559,18 @@ absl::Status LlmLiteRtCompiledModelExecutorBase::PrepareFirstPrefillAfterDecode(
     LITERT_RETURN_IF_ERROR(llm_context_->processed_context()
                                .processed_tokens()
                                .ReduceTokenCandidates(token_index_to_reduce));
-    LITERT_RETURN_IF_ERROR(
-        CopyKvCacheBuffers(output_heads, token_index_to_reduce,
-                           *input_kv_cache_buffers_, kv_cache_buffers_1_));
+    void* command_queue = nullptr;
+#if defined(__APPLE__)
+    bool has_metal_memory = std::any_of(
+        input_kv_cache_buffers_->begin(), input_kv_cache_buffers_->end(),
+        [](const auto& pair) { return pair.second.IsMetalMemory(); });
+    if (has_metal_memory) {
+      command_queue = GetMetalCommandQueue(env_);
+    }
+#endif
+    LITERT_RETURN_IF_ERROR(CopyKvCacheBuffers(
+        output_heads, token_index_to_reduce, *input_kv_cache_buffers_,
+        kv_cache_buffers_1_, command_queue));
     input_kv_cache_buffers_ = &kv_cache_buffers_1_;
     output_kv_cache_buffers_ = &kv_cache_buffers_2_;
   }
@@ -1004,9 +1030,18 @@ absl::Status LlmLiteRtCompiledModelExecutorBase::PrepareFirstDecode() {
   LITERT_RETURN_IF_ERROR(decode_kv_cache_buffers_2_.has_value());
   // Broadcast the prefill kv cache buffers to the decode kv cache buffers.
   // This is only needed when decode batch size > 1.
+  void* command_queue = nullptr;
+#if defined(__APPLE__)
+  bool has_metal_memory = std::any_of(
+      input_kv_cache_buffers_->begin(), input_kv_cache_buffers_->end(),
+      [](const auto& pair) { return pair.second.IsMetalMemory(); });
+  if (has_metal_memory) {
+    command_queue = GetMetalCommandQueue(env_);
+  }
+#endif
   LITERT_RETURN_IF_ERROR(CopyKvCacheBuffers(
       output_heads, /*src_index_to_copy_on_prefill=*/-1,
-      *input_kv_cache_buffers_, *decode_kv_cache_buffers_1_));
+      *input_kv_cache_buffers_, *decode_kv_cache_buffers_1_, command_queue));
   input_kv_cache_buffers_ = &decode_kv_cache_buffers_1_.value();
   output_kv_cache_buffers_ = &decode_kv_cache_buffers_2_.value();
 
@@ -1577,15 +1612,7 @@ absl::Status LlmLiteRtCompiledModelExecutorStatic::Prefill(
           prefill_signature, prefill_length, prefill_length,
           prefill_input_buffers_[prefill_signature]));
     }
-    // TODO(b/494284915): Switch to use async prefill for Metal backend.
-    if (!do_prefill_sync_.has_value()) {
-      do_prefill_sync_ = std::any_of(
-          prefill_input_buffers_[prefill_signature].begin(),
-          prefill_input_buffers_[prefill_signature].end(),
-          [](const auto& pair) { return pair.second.IsMetalMemory(); });
-    }
-    bool async = !*do_prefill_sync_ &&
-                 (i < work_groups.size() - 1 || !params.GetWaitForCompletion());
+    bool async = i < work_groups.size() - 1 || !params.GetWaitForCompletion();
     RETURN_IF_ERROR(PrefillInternal(
         prefill_signature, prefill_input_buffers_[prefill_signature],
         ids.subspan(/*pos=*/0, prefill_length), async));

runtime/executor/llm_litert_compiled_model_executor.h

@@ -440,7 +440,6 @@ class LlmLiteRtCompiledModelExecutorStatic
       std::string /*prefill_signature_name*/,
       absl::flat_hash_map<absl::string_view /*input_name*/, TensorBuffer>>
       prefill_input_buffers_;
-  std::optional<bool> do_prefill_sync_;
 };
 
 // The dynamic executor for the prefill-decode compiled model.

runtime/executor/llm_litert_compiled_model_executor_test.cc

@@ -789,7 +789,8 @@ absl::StatusOr<
 CreateDynamicExecutor(Environment& env, absl::string_view model_path,
                       uint32_t kv_increment_size = 8,
                       int prefill_chunk_size = -1) {
-  auto path = std::filesystem::path(::testing::SrcDir()) / model_path;
+  auto path =
+      std::filesystem::path(::testing::SrcDir()) / std::string(model_path);
   ASSIGN_OR_RETURN(auto model_resources,
                    CreateExecutorModelResourcesLitertLm(path.string()));
   ASSIGN_OR_RETURN(auto model_assets, ModelAssets::Create(path.string()));

runtime/executor/metal_utils.cc

@@ -0,0 +1,99 @@
+// Copyright 2026 The ODML Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#if defined(__APPLE__)
+#include "runtime/executor/metal_utils.h"
+
+#include <dlfcn.h>
+
+#include <cstddef>
+#include <variant>
+
+#include "absl/log/absl_log.h"  // from @com_google_absl
+#include "absl/status/statusor.h"  // from @com_google_absl
+#include "litert/cc/litert_environment_options.h"  // from @litert
+
+namespace litert::lm {
+
+typedef int (*CopyKvCacheMetalFn)(void*, void*, int, int, size_t, size_t,
+                                  void*);
+static CopyKvCacheMetalFn g_copy_kv_cache_metal = []() -> CopyKvCacheMetalFn {
+  return reinterpret_cast<CopyKvCacheMetalFn>(
+      dlsym(RTLD_DEFAULT, "LiteRtCopyKvCacheMetal"));
+}();
+
+void* GetMetalCommandQueue(litert::Environment& env) {
+  auto env_options_or = env.GetOptions();
+  if (env_options_or.HasValue()) {
+    auto queue_option_or = env_options_or.Value().GetOption(
+        litert::EnvironmentOptions::Tag::kMetalCommandQueue);
+    if (queue_option_or.HasValue()) {
+      const auto& val = queue_option_or.Value();
+      if (std::holds_alternative<const void*>(val)) {
+        return const_cast<void*>(std::get<const void*>(val));
+      } else if (std::holds_alternative<void*>(val)) {
+        return std::get<void*>(val);
+      }
+    }
+  }
+  return nullptr;
+}
+
+int CopyKvCacheMetal(void* src_buffer_ptr, void* dst_buffer_ptr,
+                     int src_index_to_copy_on_prefill, int decode_batch_size,
+                     size_t src_buffer_size, size_t dst_buffer_size,
+                     void* command_queue) {
+  if (g_copy_kv_cache_metal == nullptr) {
+    return -1;
+  }
+  return g_copy_kv_cache_metal(src_buffer_ptr, dst_buffer_ptr,
+                               src_index_to_copy_on_prefill, decode_batch_size,
+                               src_buffer_size, dst_buffer_size, command_queue);
+}
+
+absl::StatusOr<bool> TryCopyKvCacheMetal(const litert::TensorBuffer& src_buffer,
+                                         litert::TensorBuffer& dst_buffer,
+                                         int src_index_to_copy_on_prefill,
+                                         int decode_batch_size,
+                                         void* command_queue) {
+  if (!g_copy_kv_cache_metal || !command_queue || !src_buffer.IsMetalMemory() ||
+      !dst_buffer.IsMetalMemory()) {
+    return false;
+  }
+
+  auto src_metal_buffer_or = src_buffer.GetMetalBuffer();
+  auto dst_metal_buffer_or = dst_buffer.GetMetalBuffer();
+  if (!src_metal_buffer_or.HasValue() || !dst_metal_buffer_or.HasValue()) {
+    return absl::InternalError("Failed to retrieve Metal buffer handles");
+  }
+
+  auto src_size_expected = src_buffer.PackedSize();
+  auto dst_size_expected = dst_buffer.PackedSize();
+  if (!src_size_expected.HasValue() || !dst_size_expected.HasValue()) {
+    return absl::InternalError("Failed to retrieve buffer packed sizes");
+  }
+
+  int ret = CopyKvCacheMetal(
+      src_metal_buffer_or.Value(), dst_metal_buffer_or.Value(),
+      src_index_to_copy_on_prefill, decode_batch_size,
+      src_size_expected.Value(), dst_size_expected.Value(), command_queue);
+  if (ret != 0) {
+    ABSL_LOG(WARNING) << "Metal GPU KV cache copy failed with " << ret;
+    return false;
+  }
+  return true;
+}
+
+}  // namespace litert::lm
+#endif

runtime/executor/metal_utils.h

@@ -0,0 +1,42 @@
+// Copyright 2026 The ODML Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef THIRD_PARTY_ODML_LITERT_LM_RUNTIME_EXECUTOR_METAL_UTILS_H_
+#define THIRD_PARTY_ODML_LITERT_LM_RUNTIME_EXECUTOR_METAL_UTILS_H_
+
+#include <cstddef>
+
+#include "absl/status/statusor.h"  // from @com_google_absl
+#include "litert/cc/litert_environment.h"  // from @litert
+#include "litert/cc/litert_tensor_buffer.h"  // from @litert
+
+namespace litert::lm {
+
+#if defined(__APPLE__)
+// Extract Metal command queue from LiteRT Environment options.
+void* GetMetalCommandQueue(litert::Environment& env);
+
+// Attempt to copy KV cache buffers on GPU. Checks if both buffers are Metal.
+// Returns true if copy succeeded, false if skipped (not Metal or missing
+// queue), or an error status if copy failed due to internal errors.
+absl::StatusOr<bool> TryCopyKvCacheMetal(const litert::TensorBuffer& src_buffer,
+                                         litert::TensorBuffer& dst_buffer,
+                                         int src_index_to_copy_on_prefill,
+                                         int decode_batch_size,
+                                         void* command_queue);
+#endif
+
+}  // namespace litert::lm
+
+#endif  // THIRD_PARTY_ODML_LITERT_LM_RUNTIME_EXECUTOR_METAL_UTILS_H_