From 6e0110585bc22ec07a77c7fea0582ebf9e03f9f2 Mon Sep 17 00:00:00 2001 From: Fengwu Yao Date: Mon, 18 May 2026 13:44:23 -0700 Subject: [PATCH] Internal changes only. LiteRT-LM-PiperOrigin-RevId: 917405138 --- runtime/executor/BUILD | 47 +++++---- .../llm_litert_compiled_model_executor.cc | 55 ++++++++--- .../llm_litert_compiled_model_executor.h | 1 - ...llm_litert_compiled_model_executor_test.cc | 3 +- runtime/executor/metal_utils.cc | 99 +++++++++++++++++++ runtime/executor/metal_utils.h | 42 ++++++++ 6 files changed, 214 insertions(+), 33 deletions(-) create mode 100644 runtime/executor/metal_utils.cc create mode 100644 runtime/executor/metal_utils.h diff --git a/runtime/executor/BUILD b/runtime/executor/BUILD index 08d41abe5..4abb85686 100644 --- a/runtime/executor/BUILD +++ b/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( diff --git a/runtime/executor/llm_litert_compiled_model_executor.cc b/runtime/executor/llm_litert_compiled_model_executor.cc index 1ad57f91d..da8e79bfb 100644 --- a/runtime/executor/llm_litert_compiled_model_executor.cc +++ b/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 #include #include @@ -143,13 +147,26 @@ absl::Status CopyKvCacheBuffers( const absl::flat_hash_map& src_kv_cache_buffers, const absl::flat_hash_map& - 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(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)); diff --git a/runtime/executor/llm_litert_compiled_model_executor.h b/runtime/executor/llm_litert_compiled_model_executor.h index 16a27e206..51a809ad4 100644 --- a/runtime/executor/llm_litert_compiled_model_executor.h +++ b/runtime/executor/llm_litert_compiled_model_executor.h @@ -440,7 +440,6 @@ class LlmLiteRtCompiledModelExecutorStatic std::string /*prefill_signature_name*/, absl::flat_hash_map> prefill_input_buffers_; - std::optional do_prefill_sync_; }; // The dynamic executor for the prefill-decode compiled model. diff --git a/runtime/executor/llm_litert_compiled_model_executor_test.cc b/runtime/executor/llm_litert_compiled_model_executor_test.cc index 9108c2e34..c7ec7345c 100644 --- a/runtime/executor/llm_litert_compiled_model_executor_test.cc +++ b/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())); diff --git a/runtime/executor/metal_utils.cc b/runtime/executor/metal_utils.cc new file mode 100644 index 000000000..f9bc67ba3 --- /dev/null +++ b/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 + +#include +#include + +#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( + 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(val)) { + return const_cast(std::get(val)); + } else if (std::holds_alternative(val)) { + return std::get(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 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 diff --git a/runtime/executor/metal_utils.h b/runtime/executor/metal_utils.h new file mode 100644 index 000000000..a228d6a4a --- /dev/null +++ b/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 + +#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 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_