From 584fe7d16713575bab6f2fa4e693cb8913eb77ac Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Sun, 18 Oct 2026 16:46:09 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- backends/cuda/batching/cuda_executor.cpp | 234 ++++++++++++++++-- backends/cuda/batching/cuda_executor.h | 43 ++++ .../llm/batching/test/session_table_test.cpp | 48 ++++ extension/llm/batching/util/session_table.cpp | 16 +- extension/llm/batching/util/session_table.h | 3 + 5 files changed, 325 insertions(+), 19 deletions(-) diff --git a/backends/cuda/batching/cuda_executor.cpp b/backends/cuda/batching/cuda_executor.cpp index 0608b3d8466..9fd8d4a7a13 100644 --- a/backends/cuda/batching/cuda_executor.cpp +++ b/backends/cuda/batching/cuda_executor.cpp @@ -90,6 +90,48 @@ Result step_width(const MethodMeta& meta, const char* method) { return static_cast(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> 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{static_cast(sizes[1]), sizes[2]}; +} + +bool dense(const ::executorch::aten::Tensor& tensor) { + std::int64_t stride = 1; + for (int d = tensor.dim() - 1; d >= 0; --d) { + if (tensor.size(d) > 1 && tensor.strides()[d] != stride) { + return false; + } + stride *= tensor.size(d); + } + return true; +} + // The logits width `method` produces. Result logits_width(const MethodMeta& meta, const char* method) { ET_CHECK_OR_RETURN_ERROR( @@ -479,8 +521,27 @@ Result> CudaExecutor::create( max_sessions, seq_limit.value_or(0)); + std::optional 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(new CudaExecutor( + auto executor = std::unique_ptr(new CudaExecutor( std::move(module), std::move(cache), max_sessions, @@ -491,6 +552,8 @@ Result> CudaExecutor::create( min_selected_rows, has_sampler, has_argmax)); + executor->embeddings_method_ = embeddings_method; + return executor; } bool CudaExecutor::initialize() { @@ -513,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}}) { @@ -521,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; } @@ -544,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& consume) { @@ -573,6 +669,9 @@ bool CudaExecutor::execute_raw( } } reset{raw_executing_}; raw_executing_ = true; + 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; @@ -601,8 +700,9 @@ bool CudaExecutor::execute_raw( ? slice.width : std::max(selected.size(), min_selected_rows_); selected.resize(rows, selected.empty() ? 0 : selected.back()); - if (!ctl_->declare_step(std::vector( - step.seq_ids.begin() + off, step.seq_ids.begin() + off + n))) { + if (!ctl_->declare_step( + std::vector( + step.seq_ids.begin() + off, step.seq_ids.begin() + off + n))) { return false; } std::vector token_values( @@ -623,16 +723,6 @@ bool CudaExecutor::execute_raw( } ++calls_[slice.method]; const auto logits = result->at(0).toTensor(); - auto dense = [](const ::executorch::aten::Tensor& tensor) { - std::int64_t stride = 1; - for (int d = tensor.dim() - 1; d >= 0; --d) { - if (tensor.size(d) > 1 && tensor.strides()[d] != stride) { - return false; - } - stride *= tensor.size(d); - } - return true; - }; if (logits.dim() < 2 || logits.size(logits.dim() - 1) != vocab_size_ || logits.numel() != static_cast(rows) * vocab_size_ || !is_supported_logits_type(logits.scalar_type()) || !dense(logits)) { @@ -737,6 +827,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 = step->tokens; raw.positions = step->positions; raw.seq_ids = step->seq_ids; @@ -805,7 +907,8 @@ std::optional> CudaExecutor::sample_rows( for (std::size_t row = 0; row < rows; ++row) { all_greedy = all_greedy && params[row * 4] == 0.0f; } - if (all_greedy || logits.device_type() != ::executorch::aten::DeviceType::CUDA) { + if (all_greedy || + logits.device_type() != ::executorch::aten::DeviceType::CUDA) { // Host logits are sampled greedily only; the exported sampler is on CUDA. return all_greedy ? argmax_rows(logits, rows) : std::nullopt; } @@ -929,11 +1032,112 @@ OffGraphKVMetrics CudaExecutor::kv_metrics() const { std::vector CudaExecutor::method_calls() const { std::vector 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& consume) { + using ::executorch::aten::SizesType; + const auto& embeddings = *step.embeddings; + const int total = static_cast(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() || !dense(embeddings)) { + 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(method.hidden) * + ::executorch::runtime::elementSize(method.dtype); + const auto* data = + static_cast(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 selected; + std::vector 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(method.max_tokens)) { + return false; + } + const std::size_t rows = + std::max(selected.size(), min_selected_rows_); + selected.resize(rows, selected.empty() ? 0 : selected.back()); + if (!ctl_->declare_step( + std::vector( + 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 rows_bytes( + static_cast(width) * row_bytes, 0); + std::copy( + data + static_cast(off) * row_bytes, + data + static_cast(off + n) * row_bytes, + rows_bytes.begin()); + auto input = make_tensor_ptr( + {1, + static_cast(width), + static_cast(method.hidden)}, + std::move(rows_bytes), + method.dtype); + std::vector position_values( + step.positions.begin() + off, step.positions.begin() + off + n); + position_values.resize(width, 0); + auto positions = make_tensor_ptr( + {static_cast(width)}, std::move(position_values)); + auto selector = + make_tensor_ptr({static_cast(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(rows) * vocab_size_ || + !is_supported_logits_type(logits.scalar_type()) || !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()) || !dense(*taps)) { + return false; + } + } + if (!consume({off, n, logits, taps ? &*taps : nullptr, indices})) { + return false; + } + } + return true; +} + } // namespace executorch::backends::cuda::batching diff --git a/backends/cuda/batching/cuda_executor.h b/backends/cuda/batching/cuda_executor.h index b117815fbd3..ccd96afb9eb 100644 --- a/backends/cuda/batching/cuda_executor.h +++ b/backends/cuda/batching/cuda_executor.h @@ -33,6 +33,7 @@ #include #include #include +#include #include #include // 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 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 accepts; + // Host [1, rows, hidden] embeddings for every row of `inputs`, in order. + std::function< + ::executorch::runtime::Result<::executorch::extension::TensorPtr>( + const std::vector&)> + 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> sample_on_device( const ::executorch::aten::Tensor& logits, const std::vector& sessions); + bool execute_embeddings( + const RawTargetStep& step, + const std::function& 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 methods_; std::vector calls_; + struct EmbeddingsMethod { + int max_tokens; + int min_tokens; + std::int64_t hidden; + ::executorch::aten::ScalarType dtype; + std::uint64_t calls = 0; + }; + std::optional embeddings_method_; + std::optional embedder_; int min_selected_rows_; llm_batching::util::SessionTable sessions_; bool has_sampler_; diff --git a/extension/llm/batching/test/session_table_test.cpp b/extension/llm/batching/test/session_table_test.cpp index 0efcb0a36fa..079ccf6c64f 100644 --- a/extension/llm/batching/test/session_table_test.cpp +++ b/extension/llm/batching/test/session_table_test.cpp @@ -91,6 +91,54 @@ TEST_F(SessionTableTest, PacksInputsInOrderWithTheirRows) { EXPECT_EQ(step->logit_indices, std::vector({2, -1})); } +// A prompt whose rows only its executor can embed. +class OpaquePrompt final : public batching::PreparedInput { + public: + explicit OpaquePrompt(std::size_t rows) : rows_(rows) {} + const void* kind() const override { + static char kind; + return &kind; + } + std::size_t size() const override { + return rows_; + } + batching::PreparedInputPtr suffix(std::size_t) const override { + return nullptr; + } + + private: + std::size_t rows_; +}; + +TEST_F(SessionTableTest, PreparedRowsPackAsPlaceholders) { + const auto a = *table_->open(); + const auto b = *table_->open(); + batching::Input prompt; + prompt.sid = a; + prompt.payload = + batching::PreparedInputPtr(std::make_shared(4)); + prompt.offset = 1; + prompt.size = 3; + prompt.position = 0; + prompt.produce_output = true; + batching::BatchInput batch; + batch.inputs = {input(b, {5, 6}, 0), prompt}; + ASSERT_EQ(table_->pack(batch).error(), Error::InvalidArgument); + prompt.offset = 0; + batch.inputs = {input(b, {5, 6}, 0), prompt}; + const auto step = table_->pack(batch); + ASSERT_EQ(step.error(), Error::Ok); + EXPECT_TRUE(step->prepared); + EXPECT_EQ(step->tokens, std::vector({5, 6, 0, 0, 0})); + EXPECT_EQ(step->positions, std::vector({0, 1, 0, 1, 2})); + EXPECT_EQ(step->logit_indices, std::vector({1, 4})); + prompt.size = 5; + batch.inputs = {prompt}; + EXPECT_EQ(table_->pack(batch).error(), Error::InvalidArgument); + batch.inputs = {input(b, {5, 6}, 0)}; + EXPECT_FALSE(table_->pack(batch)->prepared); +} + TEST_F(SessionTableTest, ConsecutiveChunksOfOnePromptAbut) { const auto a = *table_->open(); batching::BatchInput batch; diff --git a/extension/llm/batching/util/session_table.cpp b/extension/llm/batching/util/session_table.cpp index ab2333783cb..3be64ba5f1d 100644 --- a/extension/llm/batching/util/session_table.cpp +++ b/extension/llm/batching/util/session_table.cpp @@ -195,8 +195,14 @@ Result SessionTable::pack(const BatchInput& batch) { const std::int32_t seq_id = seq_it->second.seq_id; const auto* tokens = std::get_if(&input.payload); const auto* token_input = tokens && *tokens ? tokens->get() : nullptr; - if (input.size == 0 || !token_input || input.offset > token_input->size() || - input.size > token_input->size() - input.offset) { + const auto* prepared = std::get_if(&input.payload); + const auto* prepared_input = + prepared && *prepared ? prepared->get() : nullptr; + const std::size_t available = token_input ? token_input->size() + : prepared_input ? prepared_input->size() + : 0; + if (input.size == 0 || (!token_input && !prepared_input) || + input.offset > available || input.size > available - input.offset) { ET_LOG( Error, "pack: session %" PRId64 " gave a slice its tokens do not hold", @@ -254,9 +260,11 @@ Result SessionTable::pack(const BatchInput& batch) { return Error::OutOfResources; } - const Token* slice = token_input->data() + input.offset; + const Token* slice = + token_input ? token_input->data() + input.offset : nullptr; + step.prepared = step.prepared || prepared_input != nullptr; for (std::size_t k = 0; k < input.size; ++k) { - step.tokens.push_back(static_cast(slice[k])); + step.tokens.push_back(slice ? static_cast(slice[k]) : 0); step.positions.push_back(start + static_cast(k)); } step.seq_ids.insert(step.seq_ids.end(), input.size, seq_id); diff --git a/extension/llm/batching/util/session_table.h b/extension/llm/batching/util/session_table.h index 756683302de..b61c4e2a53c 100644 --- a/extension/llm/batching/util/session_table.h +++ b/extension/llm/batching/util/session_table.h @@ -43,6 +43,9 @@ struct ET_EXPERIMENTAL PackedStep { std::vector seq_ids; // Per input: the token whose logits it wants, or -1 for none. std::vector logit_indices; + // Some input is a PreparedInput: its rows hold placeholder tokens, so the + // step must run from the executor's materialized embeddings. + bool prepared = false; }; // The logits rows a forward over tokens [offset, offset + length) must select: