1 /* 2 * Copyright (c) Meta Platforms, Inc. and affiliates. 3 * All rights reserved. 4 * 5 * This source code is licensed under the BSD-style license found in the 6 * LICENSE file in the root directory of this source tree. 7 */ 8 9 // Given a image tensor, prefill the KV cache of a multimodal LLM. 10 11 #pragma once 12 13 #include <executorch/extension/llm/runner/image.h> 14 #include <executorch/extension/module/module.h> 15 #include <executorch/runtime/platform/compiler.h> 16 17 namespace executorch { 18 namespace extension { 19 namespace llm { 20 21 // Assuming kv cache and parallel prefill are enabled. 22 class ET_EXPERIMENTAL ImagePrefiller { 23 public: ImagePrefiller(::executorch::extension::Module * module)24 explicit ImagePrefiller(::executorch::extension::Module* module) 25 : module_(module) {} 26 27 /** 28 * Prefill an LLM Module with the given image input. 29 * @param image The image input to the multimodal LLM. 30 * @param start_pos The starting position in KV cache of the input in the LLM. 31 * It's passed as reference and will be updated inside this function. 32 * @return The next token of the LLM Module after prefill. 33 */ 34 virtual ::executorch::runtime::Result<executorch::aten::Tensor> prefill( 35 Image& image, 36 int64_t& start_pos) = 0; 37 38 virtual ::executorch::runtime::Error load() = 0; 39 virtual bool is_method_loaded() = 0; 40 41 virtual ~ImagePrefiller() = default; 42 43 protected: 44 Module* module_; 45 }; 46 47 } // namespace llm 48 } // namespace extension 49 } // namespace executorch 50 51 namespace torch { 52 namespace executor { 53 // TODO(T197294990): Remove these deprecated aliases once all users have moved 54 // to the new `::executorch` namespaces. 55 using ::executorch::extension::llm::ImagePrefiller; 56 } // namespace executor 57 } // namespace torch 58