[executorch][cuda] Run CudaExecutor steps from prepared embeddings

backends/cuda/batching/cuda_executor.cpp

@@ -90,6 +90,37 @@ Result<int> step_width(const MethodMeta& meta, const char* method) {
   return static_cast<int>(token_sizes[1]);
 }
 
+// kEmbeddingsMethod's decoder input: [1, T, hidden] in a float dtype, with
+// Long[T] positions and a rank-one Long selector. Returns T's upper bound.
+Result<std::pair<int, std::int64_t>> embeddings_width(
+    const MethodMeta& meta,
+    ::executorch::aten::ScalarType* dtype) {
+  const auto embeddings = meta.input_tensor_meta(0);
+  const auto positions = meta.input_tensor_meta(1);
+  const auto selector = meta.input_tensor_meta(2);
+  ET_CHECK_OR_RETURN_ERROR(
+      meta.num_inputs() == 3 && embeddings.ok() && positions.ok() &&
+          selector.ok(),
+      InvalidProgram,
+      "CudaExecutor: %s must take embeddings, positions and a logits selector",
+      kEmbeddingsMethod);
+  const auto sizes = embeddings->sizes();
+  const auto position_sizes = positions->sizes();
+  ET_CHECK_OR_RETURN_ERROR(
+      is_supported_logits_type(embeddings->scalar_type()) &&
+          sizes.size() == 3 && sizes[0] == 1 && sizes[1] > 0 && sizes[2] > 0 &&
+          positions->scalar_type() == ::executorch::aten::ScalarType::Long &&
+          position_sizes.size() == 1 && position_sizes[0] == sizes[1] &&
+          selector->scalar_type() == ::executorch::aten::ScalarType::Long &&
+          selector->sizes().size() == 1,
+      InvalidProgram,
+      "CudaExecutor: %s must take [1, T, hidden] embeddings, Long[T] positions "
+      "and a rank-one Long selector",
+      kEmbeddingsMethod);
+  *dtype = embeddings->scalar_type();
+  return std::pair<int, std::int64_t>{static_cast<int>(sizes[1]), sizes[2]};
+}
+
 // The logits width `method` produces.
 Result<std::int64_t> logits_width(const MethodMeta& meta, const char* method) {
   ET_CHECK_OR_RETURN_ERROR(
@@ -490,8 +521,27 @@ Result<std::unique_ptr<CudaExecutor>> CudaExecutor::create(
       max_sessions,
       seq_limit.value_or(0));
 
+  std::optional<EmbeddingsMethod> embeddings_method;
+  if (names.count(kEmbeddingsMethod) != 0) {
+    ET_ASSIGN_OR_RETURN(meta, module->method_meta(kEmbeddingsMethod));
+    ::executorch::aten::ScalarType dtype{};
+    ET_ASSIGN_OR_RETURN(width, embeddings_width(meta, &dtype));
+    ET_ASSIGN_OR_RETURN(vocab, logits_width(meta, kEmbeddingsMethod));
+    ET_CHECK_OK_OR_RETURN_ERROR(backend_of(meta, kEmbeddingsMethod).error());
+    ET_CHECK_OR_RETURN_ERROR(
+        vocab == method_vocab && width.first <= max_step_tokens &&
+            min_selected_rows <= width.first,
+        InvalidProgram,
+        "CudaExecutor: %s disagrees with the token forwards",
+        kEmbeddingsMethod);
+    // Exported, like forward_others, from min_selected_rows tokens; narrower
+    // steps pad up to it.
+    embeddings_method =
+        EmbeddingsMethod{width.first, min_selected_rows, width.second, dtype};
+  }
+
   ET_CHECK_OK_OR_RETURN_ERROR(set_backend_options(options, methods));
-  return std::unique_ptr<CudaExecutor>(new CudaExecutor(
+  auto executor = std::unique_ptr<CudaExecutor>(new CudaExecutor(
       std::move(module),
       std::move(cache),
       max_sessions,
@@ -502,6 +552,8 @@ Result<std::unique_ptr<CudaExecutor>> CudaExecutor::create(
       min_selected_rows,
       has_sampler,
       has_argmax));
+  executor->embeddings_method_ = embeddings_method;
+  return executor;
 }
 
 bool CudaExecutor::initialize() {
@@ -524,6 +576,15 @@ bool CudaExecutor::initialize() {
       return false;
     }
   }
+  if (embeddings_method_ &&
+      module_->load_method(
+          kEmbeddingsMethod,
+          /*planned_memory=*/nullptr,
+          /*event_tracer=*/nullptr,
+          &options_map) != Error::Ok) {
+    ET_LOG(Error, "CudaExecutor: could not load %s", kEmbeddingsMethod);
+    return false;
+  }
   for (const auto& [present, name] :
        {std::pair{has_sampler_, kSampleMethod},
         std::pair{has_argmax_, kArgmaxMethod}}) {
@@ -532,6 +593,11 @@ bool CudaExecutor::initialize() {
       return false;
     }
   }
+  if (embedder_ && embedder_->initialize &&
+      embedder_->initialize() != Error::Ok) {
+    ET_LOG(Error, "CudaExecutor: the input embedder failed to initialize");
+    return false;
+  }
   initialized_ = true;
   return true;
 }
@@ -555,6 +621,25 @@ void CudaExecutor::set_sampling(
   sessions_.set_sampling(session, params, seed);
 }
 
+bool CudaExecutor::accepts(const llm_batching::PreparedInput& input) const {
+  return embedder_ && embedder_->accepts && embedder_->accepts(input);
+}
+
+Error CudaExecutor::set_input_embedder(InputEmbedder embedder) {
+  ET_CHECK_OR_RETURN_ERROR(
+      !initialized_ && embedder.accepts && embedder.materialize,
+      InvalidArgument,
+      "CudaExecutor: an embedder needs accepts and materialize, before "
+      "initialize()");
+  ET_CHECK_OR_RETURN_ERROR(
+      embeddings_method_.has_value(),
+      NotSupported,
+      "CudaExecutor: the program exports no %s",
+      kEmbeddingsMethod);
+  embedder_ = std::move(embedder);
+  return Error::Ok;
+}
+
 bool CudaExecutor::execute_raw(
     const RawTargetStep& step,
     const std::function<bool(const RawTargetOutput&)>& consume) {
@@ -590,6 +675,9 @@ bool CudaExecutor::execute_raw(
     }
     bool& flag;
   } reset(raw_executing_);
+  if (step.embeddings) {
+    return execute_embeddings(step, consume);
+  }
   for (const auto& slice : plan_slices(total, methods_)) {
     const int off = slice.offset;
     const int n = slice.length;
@@ -756,6 +844,18 @@ bool CudaExecutor::execute(const BatchInput& batch, BatchOutput& out) {
   }
 
   RawTargetStep raw;
+  if (step->prepared) {
+    if (!embedder_) {
+      ET_LOG(Error, "CudaExecutor: a prepared input needs an input embedder");
+      return false;
+    }
+    auto embeddings = embedder_->materialize(batch.inputs);
+    if (!embeddings.ok()) {
+      ET_LOG(Error, "CudaExecutor: could not materialize the batch's inputs");
+      return false;
+    }
+    raw.embeddings = std::move(*embeddings);
+  }
   raw.tokens = std::move(step->tokens);
   raw.positions = std::move(step->positions);
   raw.seq_ids = std::move(step->seq_ids);
@@ -954,11 +1054,112 @@ OffGraphKVMetrics CudaExecutor::kv_metrics() const {
 
 std::vector<ForwardMethodCalls> CudaExecutor::method_calls() const {
   std::vector<ForwardMethodCalls> calls;
-  calls.reserve(methods_.size());
+  calls.reserve(methods_.size() + 1);
   for (std::size_t i = 0; i < methods_.size(); ++i) {
     calls.push_back({methods_[i].name, calls_[i]});
   }
+  if (embeddings_method_) {
+    calls.push_back({kEmbeddingsMethod, embeddings_method_->calls});
+  }
   return calls;
 }
 
+bool CudaExecutor::execute_embeddings(
+    const RawTargetStep& step,
+    const std::function<bool(const RawTargetOutput&)>& consume) {
+  using ::executorch::aten::SizesType;
+  const auto& embeddings = *step.embeddings;
+  const int total = static_cast<int>(step.positions.size());
+  if (!embeddings_method_ || embeddings.dim() != 3 || embeddings.size(0) != 1 ||
+      embeddings.size(1) != total ||
+      embeddings.size(2) != embeddings_method_->hidden ||
+      embeddings.scalar_type() != embeddings_method_->dtype ||
+      !embeddings.device().is_cpu() || !is_dense(embeddings) || total <= 0 ||
+      embeddings.const_data_ptr() == nullptr) {
+    ET_LOG(
+        Error,
+        "CudaExecutor: embeddings must be host [1, %d, %" PRId64 "] tensors",
+        total,
+        embeddings_method_ ? embeddings_method_->hidden : 0);
+    return false;
+  }
+  const auto& method = *embeddings_method_;
+  const std::size_t row_bytes = static_cast<std::size_t>(method.hidden) *
+      ::executorch::runtime::elementSize(method.dtype);
+  const auto* data =
+      static_cast<const std::uint8_t*>(embeddings.const_data_ptr());
+  for (int off = 0; off < total; off += method.max_tokens) {
+    const int n = std::min(method.max_tokens, total - off);
+    const int width = std::max(n, method.min_tokens);
+    std::vector<std::int64_t> selected;
+    std::vector<std::size_t> indices;
+    for (std::size_t i = 0; i < step.logit_indices.size(); ++i) {
+      const auto index = step.logit_indices[i];
+      if (index >= off && index < off + n) {
+        selected.push_back(index - off);
+        indices.push_back(i);
+      }
+    }
+    if (selected.size() > static_cast<std::size_t>(method.max_tokens)) {
+      return false;
+    }
+    const std::size_t rows =
+        std::max<std::size_t>(selected.size(), min_selected_rows_);
+    selected.resize(rows, selected.empty() ? 0 : selected.back());
+    if (!ctl_->declare_step(std::vector<std::int32_t>(
+            step.seq_ids.begin() + off, step.seq_ids.begin() + off + n))) {
+      return false;
+    }
+    // Rows past n are padding the cache writes to its scratch cell.
+    std::vector<std::uint8_t> rows_bytes(
+        static_cast<std::size_t>(width) * row_bytes, 0);
+    std::copy(
+        data + static_cast<std::size_t>(off) * row_bytes,
+        data + static_cast<std::size_t>(off + n) * row_bytes,
+        rows_bytes.begin());
+    auto input = make_tensor_ptr(
+        {1,
+         static_cast<SizesType>(width),
+         static_cast<SizesType>(method.hidden)},
+        std::move(rows_bytes),
+        method.dtype);
+    std::vector<std::int64_t> position_values(
+        step.positions.begin() + off, step.positions.begin() + off + n);
+    position_values.resize(width, 0);
+    auto positions = make_tensor_ptr(
+        {static_cast<SizesType>(width)}, std::move(position_values));
+    auto selector =
+        make_tensor_ptr({static_cast<SizesType>(rows)}, std::move(selected));
+    auto result =
+        module_->execute(kEmbeddingsMethod, {input, positions, selector});
+    if (!result.ok() || result->empty() || !result->at(0).isTensor()) {
+      ET_LOG(Error, "CudaExecutor: %s failed", kEmbeddingsMethod);
+      return false;
+    }
+    ++embeddings_method_->calls;
+    const auto logits = result->at(0).toTensor();
+    if (logits.dim() < 2 || logits.size(logits.dim() - 1) != vocab_size_ ||
+        logits.numel() != static_cast<std::int64_t>(rows) * vocab_size_ ||
+        !is_supported_logits_type(logits.scalar_type()) || !is_dense(logits)) {
+      return false;
+    }
+    std::optional<::executorch::aten::Tensor> taps;
+    if (result->size() > 1) {
+      if (!result->at(1).isTensor()) {
+        return false;
+      }
+      taps = result->at(1).toTensor();
+      if (taps->dim() != 3 || taps->size(0) != 1 || taps->size(1) != width ||
+          taps->size(2) <= 0 ||
+          !is_supported_logits_type(taps->scalar_type()) || !is_dense(*taps)) {
+        return false;
+      }
+    }
+    if (!consume({off, n, logits, taps ? &*taps : nullptr, indices})) {
+      return false;
+    }
+  }
+  return true;
+}
+
 } // namespace executorch::backends::cuda::batching

backends/cuda/batching/cuda_executor.h

@@ -33,6 +33,7 @@
 #include <executorch/extension/llm/cache/cache.h>
 #include <executorch/extension/llm/cache/cache_registry.h>
 #include <executorch/extension/module/module.h>
+#include <executorch/extension/tensor/tensor_ptr.h>
 #include <executorch/runtime/core/result.h>
 #include <executorch/runtime/platform/compiler.h> // ET_EXPERIMENTAL
 
@@ -56,6 +57,10 @@ inline constexpr char kMinSelectedRowsMethod[] = "get_min_prefill_chunk";
 // both returning Long [R], and only the tokens come back to the host.
 inline constexpr char kSampleMethod[] = "sample";
 inline constexpr char kArgmaxMethod[] = "sample_argmax";
+// Optional dynamic forward over the same cache that takes decoder inputs
+// rather than tokens: (embeddings [1, T, hidden], Long[T] positions, Long[R]
+// selector). It runs only the steps carrying prepared inputs, e.g. images.
+inline constexpr char kEmbeddingsMethod[] = "forward_from_embeddings";
 
 // Process-wide CUDA backend options create() sets before the methods load.
 struct CudaExecutorOptions {
@@ -97,6 +102,22 @@ struct RawTargetStep {
   // Global token-row indices; -1 discards a result. Repeats and reordering are
   // allowed. Callback selected_indices maps rows back into this vector.
   std::vector<std::int64_t> logit_indices;
+  // Host [1, tokens, hidden] decoder inputs. When set, the step runs through
+  // kEmbeddingsMethod and `tokens` is not read.
+  ::executorch::extension::TensorPtr embeddings;
+};
+
+// Turns a batch holding prepared inputs into its decoder inputs, for a model
+// whose prompts carry more than tokens. Runs on the engine thread.
+struct InputEmbedder {
+  // Loads what materialize needs; called from initialize().
+  std::function<::executorch::runtime::Error()> initialize;
+  std::function<bool(const llm_batching::PreparedInput&)> accepts;
+  // Host [1, rows, hidden] embeddings for every row of `inputs`, in order.
+  std::function<
+      ::executorch::runtime::Result<::executorch::extension::TensorPtr>(
+          const std::vector<llm_batching::Input>&)>
+      materialize;
 };
 
 struct RawTargetOutput {
@@ -162,6 +183,16 @@ class ET_EXPERIMENTAL CudaExecutor : public llm_batching::Executor {
   bool execute(
       const llm_batching::BatchInput& batch,
       llm_batching::BatchOutput& out) override;
+  // Whether the installed embedder takes `input`; false without one.
+  bool accepts(const llm_batching::PreparedInput& input) const override;
+
+  // Installs the embedder for prepared inputs, before initialize(). The program
+  // must export kEmbeddingsMethod.
+  ::executorch::runtime::Error set_input_embedder(InputEmbedder embedder);
+  // The decoder input width kEmbeddingsMethod takes, or 0 without it.
+  std::int64_t embeddings_hidden_size() const {
+    return embeddings_method_ ? embeddings_method_->hidden : 0;
+  }
 
   // Engine-thread-only raw access for model adapters. Output storage is
   // borrowed until the callback returns: copy/consume every needed output
@@ -230,6 +261,9 @@ class ET_EXPERIMENTAL CudaExecutor : public llm_batching::Executor {
   std::optional<std::vector<llm_batching::Token>> sample_on_device(
       const ::executorch::aten::Tensor& logits,
       const std::vector<llm_batching::SessionId>& sessions);
+  bool execute_embeddings(
+      const RawTargetStep& step,
+      const std::function<bool(const RawTargetOutput&)>& consume);
 
   // Ordered so the module dies first, releasing the delegates that resolved
   // the cache before the registry entry naming it goes.
@@ -243,6 +277,15 @@ class ET_EXPERIMENTAL CudaExecutor : public llm_batching::Executor {
   // Sorted by max_tokens, ascending; the dynamic method, if any, is last.
   std::vector<ForwardMethod> methods_;
   std::vector<std::uint64_t> calls_;
+  struct EmbeddingsMethod {
+    int max_tokens;
+    int min_tokens;
+    std::int64_t hidden;
+    ::executorch::aten::ScalarType dtype;
+    std::uint64_t calls = 0;
+  };
+  std::optional<EmbeddingsMethod> embeddings_method_;
+  std::optional<InputEmbedder> embedder_;
   int min_selected_rows_;
   llm_batching::util::SessionTable sessions_;
   bool has_sampler_;