NPUW: Batched scoring as a composable element (reranker / embedding) #36375
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,213 @@ | ||
| // Copyright (C) 2018-2026 Intel Corporation | ||
| // SPDX-License-Identifier: Apache-2.0 | ||
| // | ||
|
|
||
| #include "batched.hpp" | ||
|
|
||
| #include <algorithm> | ||
| #include <utility> | ||
|
|
||
| #include "../../logging.hpp" | ||
| #include "../../util.hpp" | ||
| #include "intel_npu/npuw_private_properties.hpp" | ||
| #include "openvino/core/except.hpp" | ||
| #include "openvino/runtime/make_tensor.hpp" | ||
| #include "openvino/runtime/tensor.hpp" | ||
|
|
||
| bool ov::npuw::batched::requested(const ov::AnyMap& properties) { | ||
| const auto is_set = [&properties](const std::string& key) { | ||
| const auto it = properties.find(key); | ||
| return it != properties.end() && it->second.as<bool>(); | ||
| }; | ||
| return is_set(ov::intel_npu::npuw::text_rerank::enabled.name()) || | ||
| is_set(ov::intel_npu::npuw::text_embed::enabled.name()); | ||
| } | ||
|
|
||
| std::shared_ptr<ov::npuw::ICompiledModel> ov::npuw::batched::CompiledModel::create( | ||
| const std::shared_ptr<ov::npuw::ICompiledModel>& inner, | ||
| const std::shared_ptr<const ov::IPlugin>& plugin, | ||
| bool enabled) { | ||
| OPENVINO_ASSERT(inner != nullptr, "Batched compiled model requires an inner compiled model"); | ||
|
|
||
| // No-op wrapper: hand back the inner model unchanged for the zero-overhead path. | ||
| if (!enabled) { | ||
| return inner; | ||
| } | ||
| return std::make_shared<CompiledModel>(inner, plugin); | ||
| } | ||
|
|
||
| ov::npuw::batched::CompiledModel::CompiledModel(const std::shared_ptr<ov::npuw::ICompiledModel>& inner, | ||
| const std::shared_ptr<const ov::IPlugin>& plugin) | ||
| : ov::npuw::ICompiledModel(nullptr, plugin), // I/O comes from the inner via inputs()/outputs() | ||
| m_inner(inner) { | ||
| OPENVINO_ASSERT(m_inner != nullptr, "Batched compiled model requires an inner compiled model"); | ||
| } | ||
|
|
||
| const std::vector<ov::Output<const ov::Node>>& ov::npuw::batched::CompiledModel::inputs() const { | ||
| return m_inner->inputs(); | ||
| } | ||
|
|
||
| const std::vector<ov::Output<const ov::Node>>& ov::npuw::batched::CompiledModel::outputs() const { | ||
| return m_inner->outputs(); | ||
| } | ||
|
|
||
| void ov::npuw::batched::CompiledModel::export_model(std::ostream& model) const { | ||
| // The element is a runtime-only decorator: the blob is the inner's blob, and the | ||
| // entry points re-apply the wrapper on import based on the properties. | ||
| m_inner->export_model(model); | ||
|
Contributor
There was a problem hiding this comment. So... I don't think it would actually work. Even though it is just a runtime wrapper you must recreate it from somewhere to create it when you would deserialize. And I don't see where it is happens. Need to figure out how to handle that properly. |
||
| } | ||
|
|
||
| std::shared_ptr<const ov::Model> ov::npuw::batched::CompiledModel::get_runtime_model() const { | ||
| return m_inner->get_runtime_model(); | ||
| } | ||
|
|
||
| void ov::npuw::batched::CompiledModel::set_property(const ov::AnyMap& properties) { | ||
| m_inner->set_property(properties); | ||
| } | ||
|
|
||
| ov::Any ov::npuw::batched::CompiledModel::get_property(const std::string& name) const { | ||
| return m_inner->get_property(name); | ||
| } | ||
|
|
||
| void ov::npuw::batched::CompiledModel::release_memory() { | ||
| m_inner->release_memory(); | ||
| } | ||
|
|
||
| std::shared_ptr<ov::ISyncInferRequest> ov::npuw::batched::CompiledModel::create_sync_infer_request() const { | ||
| auto self = std::static_pointer_cast<const ov::ICompiledModel>(shared_from_this()); | ||
| auto inner_request = m_inner->create_infer_request(); | ||
| OPENVINO_ASSERT(inner_request != nullptr, "Batched element: inner compiled model returned a null request"); | ||
| return std::make_shared<InferRequest>(self, std::move(inner_request)); | ||
| } | ||
|
|
||
| ov::npuw::batched::InferRequest::InferRequest(const std::shared_ptr<const ov::ICompiledModel>& compiled_model, | ||
| std::shared_ptr<ov::IAsyncInferRequest> inner_request) | ||
| : ov::ISyncInferRequest(compiled_model), | ||
| m_inner(std::move(inner_request)) { | ||
| OPENVINO_ASSERT(m_inner != nullptr, "Batched element requires a non-null inner request"); | ||
|
|
||
| m_profile.report_on_die = ov::npuw::profiling_enabled(); | ||
| m_profile.area = "batched/execution"; | ||
|
|
||
| // Surface the inner request's own tensors as the public defaults (the ports are | ||
| // the same objects, see CompiledModel::inputs()). Nothing is allocated here: a | ||
| // batch-1 caller works directly on the inner's tensors, and a batched caller | ||
| // replaces them with its [N, ...] tensors via set_tensor(). | ||
| for (const auto& port : get_inputs()) { | ||
| if (auto tensor = m_inner->get_tensor(port)) { | ||
| set_tensor(port, tensor); | ||
| } | ||
| } | ||
| } | ||
|
|
||
| ov::npuw::batched::InferRequest::BatchedInputs ov::npuw::batched::InferRequest::extract_batch() const { | ||
| const auto& in_ports = get_inputs(); | ||
| OPENVINO_ASSERT(!in_ports.empty(), "Batched element: the wrapped model has no inputs"); | ||
|
|
||
| BatchedInputs inputs; | ||
| inputs.tensors.reserve(in_ports.size()); | ||
| for (const auto& port : in_ports) { | ||
| auto tensor = get_tensor(port); | ||
| OPENVINO_ASSERT(tensor, "Batched element: no tensor is set for input '", port.get_any_name(), "'"); | ||
| const auto& shape = tensor->get_shape(); | ||
| OPENVINO_ASSERT(!shape.empty(), | ||
| "Batched element: input '", | ||
| port.get_any_name(), | ||
| "' has no leading (batch) dimension"); | ||
| OPENVINO_ASSERT(shape[0] > 0, | ||
| "Batched element: input '", | ||
| port.get_any_name(), | ||
| "' has a zero-sized batch dimension - batch size must be > 0"); | ||
| inputs.batch = std::max(inputs.batch, shape[0]); | ||
| inputs.tensors.push_back(std::move(tensor)); | ||
| } | ||
| for (std::size_t i = 0; i < in_ports.size(); ++i) { | ||
| const std::size_t in_batch = inputs.tensors[i]->get_shape()[0]; | ||
| OPENVINO_ASSERT(in_batch == inputs.batch || in_batch == 1, | ||
| "Batched element: input '", | ||
| in_ports[i].get_any_name(), | ||
| "' has batch dimension ", | ||
| in_batch, | ||
| " which is neither the inferred batch size ", | ||
| inputs.batch, | ||
| " nor 1 (shared)."); | ||
| } | ||
| return inputs; | ||
| } | ||
|
|
||
| void ov::npuw::batched::InferRequest::infer() { | ||
| std::lock_guard<std::mutex> lock(m_mutex); | ||
|
|
||
| const auto& in_ports = get_inputs(); | ||
| const auto& out_ports = get_outputs(); | ||
|
|
||
| BatchedInputs inputs; | ||
| m_profile["1.extract_batch"].record([&]() { | ||
| inputs = extract_batch(); | ||
| }); | ||
| const std::size_t batch = inputs.batch; | ||
|
|
||
| // Unroll row by row: reset the inner variable state so each row is scored as an | ||
| // independent prompt, bind the row's [1, ...] view of every batched input, run | ||
| // the batch-1 inner request, and write the row's outputs into row `row` of the | ||
| // [N, ...] public output tensors. | ||
| const auto inner_states = m_inner->query_state(); | ||
| for (std::size_t row = 0; row < batch; ++row) { | ||
| m_profile["2.bind_row"].record([&]() { | ||
| for (const auto& state : inner_states) { | ||
| state->reset(); | ||
| } | ||
| for (std::size_t i = 0; i < in_ports.size(); ++i) { | ||
| const auto& full = inputs.tensors[i]; | ||
| m_inner->set_tensor(in_ports[i], | ||
| full->get_shape()[0] == 1 ? full : ov::npuw::util::view(full, 0, row, 1)); | ||
| } | ||
| }); | ||
| m_profile["3.inner_infer"].record([&]() { | ||
| m_inner->infer(); | ||
| }); | ||
|
|
||
| if (row == 0) { | ||
| // The wrapped model's ports are dynamic - the output shapes are only | ||
| // known once the first row has been scored. | ||
| ensure_batched_outputs(batch); | ||
| } | ||
| m_profile["4.copy_row_out"].record([&]() { | ||
| for (const auto& port : out_ports) { | ||
| m_inner->get_tensor(port)->copy_to(ov::npuw::util::view(get_tensor(port), 0, row, 1)._ptr); | ||
| } | ||
| }); | ||
| } | ||
| } | ||
|
|
||
| void ov::npuw::batched::InferRequest::ensure_batched_outputs(std::size_t batch) { | ||
| for (const auto& port : get_outputs()) { | ||
| const auto inner_out = m_inner->get_tensor(port); | ||
| OPENVINO_ASSERT(inner_out && !inner_out->get_shape().empty() && inner_out->get_shape()[0] == 1, | ||
| "Batched element: output '", | ||
| port.get_any_name(), | ||
| "' of the inner request is not a [1, ...] tensor"); | ||
| ov::Shape shape = inner_out->get_shape(); | ||
| shape[0] = batch; | ||
| const auto current = get_tensor(port); | ||
| if (!current || current->get_element_type() != inner_out->get_element_type() || current->get_shape() != shape) { | ||
| set_tensor(port, ov::get_tensor_impl(ov::Tensor(inner_out->get_element_type(), shape))); | ||
| } | ||
| } | ||
| } | ||
|
|
||
| void ov::npuw::batched::InferRequest::check_tensors() const { | ||
| // No-op: the public outputs are late-bound (allocated on infer once the batch is | ||
| // known), and the batched inputs are validated and then unrolled by infer() -- the | ||
| // per-row [1, ...] tensors are checked by the inner request. | ||
| } | ||
|
|
||
| std::vector<ov::SoPtr<ov::IVariableState>> ov::npuw::batched::InferRequest::query_state() const { | ||
| // The batched element resets inner state between rows and exposes no | ||
| // cross-call state of its own, so it presents an empty state list. | ||
| return {}; | ||
| } | ||
|
|
||
| std::vector<ov::ProfilingInfo> ov::npuw::batched::InferRequest::get_profiling_info() const { | ||
| return m_inner->get_profiling_info(); | ||
| } | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,127 @@ | ||
| // Copyright (C) 2018-2026 Intel Corporation | ||
| // SPDX-License-Identifier: Apache-2.0 | ||
| // | ||
|
|
||
| #pragma once | ||
|
|
||
| #include <memory> | ||
| #include <mutex> | ||
| #include <vector> | ||
|
|
||
| #include "../../compiled_model.hpp" | ||
| #include "../../perf.hpp" | ||
| #include "openvino/runtime/iasync_infer_request.hpp" | ||
| #include "openvino/runtime/isync_infer_request.hpp" | ||
| #include "openvino/runtime/so_ptr.hpp" | ||
|
|
||
| namespace ov::npuw::batched { | ||
|
|
||
| // True when the compile/import properties opt into batched single-shot scoring -- | ||
| // currently the text rerank and text embedding pipelines. Used by the entry points | ||
| // (npuw::ICompiledModel::create and the NPUW import path) to decide whether to | ||
| // apply the batched element. | ||
| bool requested(const ov::AnyMap& properties); | ||
|
|
||
| // A compiled-model decorator that adds batched (batch > 1) execution on top of an | ||
| // inner compiled model that only supports a batch size of 1 (as NPUW's LLM | ||
| // pipeline does, since it reshapes every sub-model to a static batch of 1). | ||
| // | ||
| // Like the other v1/elements wrappers (failsafe, accuracy_checked) it is a | ||
| // transparent decorator that exposes the inner model's I/O (inputs()/outputs() | ||
| // forward to the inner, so the ports are literally the same objects) and forwards | ||
| // everything but inference to the inner model. Unlike them it is a *fan-out* | ||
| // element: a single [N, ...] inference is unrolled into N independent [1, ...] | ||
| // inferences on the inner request, the inner variable state (KV-cache) is reset | ||
| // between rows, and the per-row outputs are written into rows of the [N, ...] | ||
| // public output tensors. | ||
| // | ||
| // This is correct for single-shot scoring workloads whose rows are independent | ||
| // -- text reranking and text embedding -- where batched and per-row results are | ||
| // identical and batching is purely a throughput/ergonomics choice. It is NOT | ||
| // valid for autoregressive generation, where state must persist across calls. | ||
| // | ||
| // It is instantiated at the NPUW entry points only: npuw::ICompiledModel::create() | ||
| // on compilation and the plugin's NPUW import path on blob import (the element is | ||
| // runtime-only and is not part of the serialized blob). Both wrap the | ||
| // LLMCompiledModel produced there from the very same model, which is what | ||
| // guarantees the inner is the matching batch-1 compilation. | ||
| class CompiledModel final : public ov::npuw::ICompiledModel { | ||
| public: | ||
| // Factory method. Returns inner unwrapped when enabled == false, keeping the | ||
| // default path zero-overhead. | ||
| static std::shared_ptr<ov::npuw::ICompiledModel> create(const std::shared_ptr<ov::npuw::ICompiledModel>& inner, | ||
| const std::shared_ptr<const ov::IPlugin>& plugin, | ||
| bool enabled); | ||
|
|
||
| CompiledModel(const std::shared_ptr<ov::npuw::ICompiledModel>& inner, | ||
| const std::shared_ptr<const ov::IPlugin>& plugin); | ||
|
|
||
| // The wrapper adds no I/O of its own -- it exposes the inner model's ports. | ||
| const std::vector<ov::Output<const ov::Node>>& inputs() const override; | ||
| const std::vector<ov::Output<const ov::Node>>& outputs() const override; | ||
|
|
||
| void export_model(std::ostream& model) const override; | ||
| std::shared_ptr<const ov::Model> get_runtime_model() const override; | ||
|
|
||
| void set_property(const ov::AnyMap& properties) override; | ||
| ov::Any get_property(const std::string& name) const override; | ||
|
|
||
| void release_memory() override; | ||
|
|
||
| std::shared_ptr<ov::ISyncInferRequest> create_sync_infer_request() const override; | ||
|
|
||
| private: | ||
| std::shared_ptr<ov::npuw::ICompiledModel> m_inner; | ||
| }; | ||
|
|
||
| // Sync infer request that unrolls a batched inference over the single-sequence | ||
| // inner request. | ||
| // | ||
| // The public input tensors default to the inner request's own tensors (surfaced in | ||
| // the constructor), so a plain batch-1 infer works exactly as on the inner. When | ||
| // the caller binds [N, ...] inputs, infer() takes N as the largest leading | ||
| // dimension across the inputs (an input with a leading dim of 1 is shared across | ||
| // rows), and for each row: resets the inner variable state, binds the row's | ||
| // [1, ...] view of every batched input, runs the inner request, and copies the | ||
| // inner outputs into row i of the [N, ...] public output tensors. Caller-bound | ||
| // output tensors are reused when they already have the right shape and type. | ||
| class InferRequest final : public ov::ISyncInferRequest { | ||
| public: | ||
| InferRequest(const std::shared_ptr<const ov::ICompiledModel>& compiled_model, | ||
| std::shared_ptr<ov::IAsyncInferRequest> inner_request); | ||
|
|
||
| void infer() override; | ||
| void check_tensors() const override; | ||
|
|
||
| std::vector<ov::SoPtr<ov::IVariableState>> query_state() const override; | ||
| std::vector<ov::ProfilingInfo> get_profiling_info() const override; | ||
|
|
||
| private: | ||
| // The public input tensors snapshotted for one infer() call, together with the | ||
| // batch size derived from them. | ||
| struct BatchedInputs { | ||
| std::vector<ov::SoPtr<ov::ITensor>> tensors; // parallel to get_inputs() | ||
| std::size_t batch = 1; | ||
| }; | ||
|
|
||
| // Snapshot the public inputs and derive the batch size: the largest leading | ||
| // dimension across them. Every input must either carry the batch ([N, ...], | ||
| // sliced per row by infer()) or be shared across rows ([1, ...], bound whole); | ||
| // anything else throws. | ||
| BatchedInputs extract_batch() const; | ||
|
|
||
| // Make the public output tensors [batch, ...] copies of the inner's [1, ...] | ||
| // outputs, reusing caller-bound tensors that already fit. The wrapped model's | ||
| // ports are dynamic, so this can only run once the first row has been scored | ||
| // and the inner output shapes are known. | ||
| void ensure_batched_outputs(std::size_t batch); | ||
|
|
||
| std::shared_ptr<ov::IAsyncInferRequest> m_inner; | ||
| mutable std::mutex m_mutex; | ||
|
|
||
| // Per-phase timings of the unroll. | ||
| using MS = ov::npuw::perf::metric<ov::npuw::perf::MSec>; | ||
| ov::npuw::perf::Profile<MS> m_profile; | ||
| }; | ||
|
|
||
| } // namespace ov::npuw::batched |
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
enabledas a parameter feels weird. Since it is passed it can be handled from outside more explicitly.