diff --git a/src/plugins/intel_npu/src/al/include/intel_npu/config/npuw_option_defs.inc b/src/plugins/intel_npu/src/al/include/intel_npu/config/npuw_option_defs.inc index 081be721abc4b8..81e187efd1d138 100644 --- a/src/plugins/intel_npu/src/al/include/intel_npu/config/npuw_option_defs.inc +++ b/src/plugins/intel_npu/src/al/include/intel_npu/config/npuw_option_defs.inc @@ -87,6 +87,7 @@ INTEL_NPU_NPUW_SIMPLE_OPT(NPUW_WHISPER_EOS_TOKEN, uint64_t, 50257, ov::intel_npu INTEL_NPU_NPUW_SIMPLE_OPT(NPUW_WHISPER_DECOMPOSE_SDPA, bool, false, ov::intel_npu::npuw::whisper, whisper_decompose_sdpa, "NPUW_WHISPER_DECOMPOSE_SDPA", LLM, EXPOSED, CACHED, ALL) INTEL_NPU_NPUW_SIMPLE_OPT(NPUW_EAGLE, bool, false, ov::intel_npu::npuw::eagle, enabled, "NPUW_EAGLE", LLM, EXPOSED, CACHED, ALL) INTEL_NPU_NPUW_SIMPLE_OPT(NPUW_TEXT_EMBED, bool, false, ov::intel_npu::npuw::text_embed, enabled, "NPUW_TEXT_EMBED", LLM, EXPOSED, CACHED, ALL) +INTEL_NPU_NPUW_SIMPLE_OPT(NPUW_TEXT_RERANK, bool, false, ov::intel_npu::npuw::text_rerank, enabled, "NPUW_TEXT_RERANK", LLM, EXPOSED, CACHED, ALL) INTEL_NPU_NPUW_SIMPLE_OPT(NPUW_KOKORO, bool, false, ov::intel_npu::npuw::kokoro, enabled, "NPUW_KOKORO", KOKORO, EXPOSED, UNCACHED, ALL) INTEL_NPU_NPUW_SIMPLE_OPT(NPUW_KOKORO_BLOCK_SIZE, uint64_t, 200, ov::intel_npu::npuw::kokoro, block_size, "NPUW_KOKORO_BLOCK_SIZE", KOKORO, EXPOSED, UNCACHED, ALL) INTEL_NPU_NPUW_SIMPLE_OPT(NPUW_KOKORO_OVERLAP_SIZE, uint64_t, 20, ov::intel_npu::npuw::kokoro, overlap_size, "NPUW_KOKORO_OVERLAP_SIZE", KOKORO, EXPOSED, UNCACHED, ALL) diff --git a/src/plugins/intel_npu/src/plugin/npuw/compiled_model.cpp b/src/plugins/intel_npu/src/plugin/npuw/compiled_model.cpp index 66c603c51054e5..ab27f734797200 100644 --- a/src/plugins/intel_npu/src/plugin/npuw/compiled_model.cpp +++ b/src/plugins/intel_npu/src/plugin/npuw/compiled_model.cpp @@ -30,6 +30,7 @@ #include "unfold_sync_infer_request.hpp" #include "util.hpp" #include "v1/elements/accuracy_checked.hpp" +#include "v1/elements/batched.hpp" #include "v1/elements/failsafe.hpp" // required for get_properties_per_device() @@ -311,7 +312,14 @@ std::shared_ptr ov::npuw::ICompiledModel::create( compiled_model = std::make_shared(model, plugin, config); } else if (properties.count(use_llm_key) && properties.at(use_llm_key).as() == true) { LOG_INFO("ov::npuw::LLMCompiledModel will be created."); - compiled_model = std::make_shared(model, plugin, config); + auto llm_compiled_model = std::make_shared(model, plugin, config); + // Single-shot scoring pipelines (text rerank / embedding) may submit batched + // [N, ...] inputs, while the LLM pipeline pins everything to a static batch + // of 1. Wrap the compiled model with the batched element, which unrolls such + // an infer row by row over the unchanged batch-1 inner request. + compiled_model = ov::npuw::batched::CompiledModel::create(llm_compiled_model, + plugin, + ov::npuw::batched::requested(properties)); } else if (properties.count(use_kokoro_key) && properties.at(use_kokoro_key).as() == true) { LOG_INFO("ov::npuw::KokoroCompiledModel will be created."); compiled_model = std::make_shared(model, plugin, config); diff --git a/src/plugins/intel_npu/src/plugin/npuw/llm_compiled_model.cpp b/src/plugins/intel_npu/src/plugin/npuw/llm_compiled_model.cpp index 53b9059d03f7a5..a23d073508c50b 100644 --- a/src/plugins/intel_npu/src/plugin/npuw/llm_compiled_model.cpp +++ b/src/plugins/intel_npu/src/plugin/npuw/llm_compiled_model.cpp @@ -815,6 +815,11 @@ ov::npuw::LLMCompiledModel::LLMCompiledModel(const std::shared_ptr& m auto use_text_embed_key = pop_option(other_props, std::string("NPUW_TEXT_EMBED")); m_is_embedding = use_text_embed_key.value_or(false).as() == true; + // NPUW_TEXT_RERANK is consumed at the entry point (npuw::ICompiledModel::create + // wraps this model with the batched element); pop it here so it doesn't leak + // into the submodel configs. + pop_option(other_props, std::string("NPUW_TEXT_RERANK")); + if (m_is_embedding) { LOG_DEBUG("Text-embedding model rebuild"); ov::npuw::util::PrepareTextEmbeddingModel(seq_len_dim).run_on_model(kvcache_model); @@ -1675,6 +1680,7 @@ void ov::npuw::LLMCompiledModel::implement_properties() { BIND(npuw::whisper::whisper_eos_token, NPUW_WHISPER_EOS_TOKEN, get), BIND(npuw::whisper::whisper_decompose_sdpa, NPUW_WHISPER_DECOMPOSE_SDPA, get), BIND(npuw::eagle::enabled, NPUW_EAGLE, get), - BIND(npuw::text_embed::enabled, NPUW_TEXT_EMBED, get)}); + BIND(npuw::text_embed::enabled, NPUW_TEXT_EMBED, get), + BIND(npuw::text_rerank::enabled, NPUW_TEXT_RERANK, get)}); #undef BIND } diff --git a/src/plugins/intel_npu/src/plugin/npuw/v1/elements/batched.cpp b/src/plugins/intel_npu/src/plugin/npuw/v1/elements/batched.cpp new file mode 100644 index 00000000000000..8d52122eefdf83 --- /dev/null +++ b/src/plugins/intel_npu/src/plugin/npuw/v1/elements/batched.cpp @@ -0,0 +1,213 @@ +// Copyright (C) 2018-2026 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#include "batched.hpp" + +#include +#include + +#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(); + }; + 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::batched::CompiledModel::create( + const std::shared_ptr& inner, + const std::shared_ptr& 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(inner, plugin); +} + +ov::npuw::batched::CompiledModel::CompiledModel(const std::shared_ptr& inner, + const std::shared_ptr& 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::npuw::batched::CompiledModel::inputs() const { + return m_inner->inputs(); +} + +const std::vector>& 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); +} + +std::shared_ptr 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::npuw::batched::CompiledModel::create_sync_infer_request() const { + auto self = std::static_pointer_cast(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(self, std::move(inner_request)); +} + +ov::npuw::batched::InferRequest::InferRequest(const std::shared_ptr& compiled_model, + std::shared_ptr 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 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::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::npuw::batched::InferRequest::get_profiling_info() const { + return m_inner->get_profiling_info(); +} diff --git a/src/plugins/intel_npu/src/plugin/npuw/v1/elements/batched.hpp b/src/plugins/intel_npu/src/plugin/npuw/v1/elements/batched.hpp new file mode 100644 index 00000000000000..f1cbec31588a4e --- /dev/null +++ b/src/plugins/intel_npu/src/plugin/npuw/v1/elements/batched.hpp @@ -0,0 +1,127 @@ +// Copyright (C) 2018-2026 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#pragma once + +#include +#include +#include + +#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 create(const std::shared_ptr& inner, + const std::shared_ptr& plugin, + bool enabled); + + CompiledModel(const std::shared_ptr& inner, + const std::shared_ptr& plugin); + + // The wrapper adds no I/O of its own -- it exposes the inner model's ports. + const std::vector>& inputs() const override; + const std::vector>& outputs() const override; + + void export_model(std::ostream& model) const override; + std::shared_ptr 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 create_sync_infer_request() const override; + +private: + std::shared_ptr 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& compiled_model, + std::shared_ptr inner_request); + + void infer() override; + void check_tensors() const override; + + std::vector> query_state() const override; + std::vector 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> 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 m_inner; + mutable std::mutex m_mutex; + + // Per-phase timings of the unroll. + using MS = ov::npuw::perf::metric; + ov::npuw::perf::Profile m_profile; +}; + +} // namespace ov::npuw::batched diff --git a/src/plugins/intel_npu/src/plugin/src/plugin.cpp b/src/plugins/intel_npu/src/plugin/src/plugin.cpp index fcc4a43dd2fb0d..9bbe8090e11885 100644 --- a/src/plugins/intel_npu/src/plugin/src/plugin.cpp +++ b/src/plugins/intel_npu/src/plugin/src/plugin.cpp @@ -24,6 +24,7 @@ #include "npuw/llm_compiled_model.hpp" #include "npuw/orc/schema_npuw.hpp" #include "npuw/serialization.hpp" +#include "npuw/v1/elements/batched.hpp" #include "openvino/core/rt_info/weightless_caching_attributes.hpp" #include "openvino/op/constant.hpp" #include "openvino/op/parameter.hpp" @@ -173,7 +174,13 @@ std::shared_ptr import_model_npuw(std::istream& stream, return ov::npuw::GQACompiledModel::import_model(stream, pluginSO, properties); } else if (compiled_model_indicator == NPUW_LLM_COMPILED_MODEL_INDICATOR) { // Properties are required for ov::weights_path - return ov::npuw::LLMCompiledModel::import_model(stream, pluginSO, properties); + auto llm_compiled_model = ov::npuw::LLMCompiledModel::import_model(stream, pluginSO, properties); + // The batched-scoring element is a runtime-only decorator and is not + // part of the blob - re-apply it from the import properties, mirroring + // ov::npuw::ICompiledModel::create(). + return ov::npuw::batched::CompiledModel::create(llm_compiled_model, + pluginSO, + ov::npuw::batched::requested(properties)); } else if (compiled_model_indicator == NPUW_COMPILED_MODEL_INDICATOR) { OPENVINO_THROW("Legacy flat NPUW CompiledModel blobs are no longer supported. Re-export the model with " "the current ORC serializer."); diff --git a/src/plugins/intel_npu/tests/unit/CMakeLists.txt b/src/plugins/intel_npu/tests/unit/CMakeLists.txt index 9dc2cf8e2f0047..123e8315db11d8 100644 --- a/src/plugins/intel_npu/tests/unit/CMakeLists.txt +++ b/src/plugins/intel_npu/tests/unit/CMakeLists.txt @@ -69,6 +69,7 @@ ov_add_test_target( ${OpenVINO_SOURCE_DIR}/src/plugins/intel_npu/src/plugin/npuw/util.cpp ${OpenVINO_SOURCE_DIR}/src/plugins/intel_npu/src/plugin/npuw/util_xarch.cpp ${OpenVINO_SOURCE_DIR}/src/plugins/intel_npu/src/plugin/npuw/v1/elements/accuracy_checked.cpp + ${OpenVINO_SOURCE_DIR}/src/plugins/intel_npu/src/plugin/npuw/v1/elements/batched.cpp ${OpenVINO_SOURCE_DIR}/src/plugins/intel_npu/src/plugin/npuw/v1/elements/failsafe.cpp ${OpenVINO_SOURCE_DIR}/src/plugins/intel_npu/src/plugin/npuw/lazy_tensor.cpp ${OpenVINO_SOURCE_DIR}/src/plugins/intel_npu/src/plugin/npuw/npuw_transformations/split_kvcache_into_blocks.cpp @@ -147,6 +148,7 @@ target_sources(${TARGET_NAME} PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/npu/executor_test.cpp ${CMAKE_CURRENT_SOURCE_DIR}/npuw/llm_compiled_model_graph_options_test.cpp ${CMAKE_CURRENT_SOURCE_DIR}/npuw/gqa_compiled_model_test.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/npuw/batched_element_test.cpp ${CMAKE_CURRENT_SOURCE_DIR}/npuw/stored_tokens_state_test.cpp ${CMAKE_CURRENT_SOURCE_DIR}/npuw/llm_infer_request_variant_switch_test.cpp ${CMAKE_CURRENT_SOURCE_DIR}/npuw/lincache_utils_test.cpp diff --git a/src/plugins/intel_npu/tests/unit/npuw/batched_element_test.cpp b/src/plugins/intel_npu/tests/unit/npuw/batched_element_test.cpp new file mode 100644 index 00000000000000..8a51ae4c91a00e --- /dev/null +++ b/src/plugins/intel_npu/tests/unit/npuw/batched_element_test.cpp @@ -0,0 +1,406 @@ +// Copyright (C) 2018-2026 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +// Unit tests for the batched v1 element (ov::npuw::batched), a model-agnostic fan-out +// decorator that unrolls a batched [N, ...] request into N batch-1 inner inferences, +// resets inner variable state between rows, and writes the per-row outputs into rows +// of the [N, ...] public output tensors. The element itself is the unit under test; +// only the inner compiled model is mocked, which keeps the tests deviceless and lets +// them observe the per-row infer/reset. It is driven through a synthetic Qwen-style +// reranker (realistic multi-input, stateful signature), and the mock echoes each +// row's first token so a test can assert output row r reflects input row r. + +#include + +#include +#include +#include +#include +#include +#include + +#include "llm_test_helpers.hpp" +#include "openvino/core/model.hpp" +#include "openvino/op/constant.hpp" +#include "openvino/op/convert.hpp" +#include "openvino/op/multiply.hpp" +#include "openvino/op/parameter.hpp" +#include "openvino/runtime/iasync_infer_request.hpp" +#include "openvino/runtime/icompiled_model.hpp" +#include "openvino/runtime/isync_infer_request.hpp" +#include "openvino/runtime/ivariable_state.hpp" +#include "openvino/runtime/make_tensor.hpp" +#include "v1/elements/batched.hpp" + +namespace { + +using ov::test::npuw::build_reranker_test_model; +using ov::test::npuw::NullPlugin; + +// Per-output value offset the mock applies so multi-output tests can tell stacked outputs apart. +constexpr float kOutputChannelOffset = 1000.0f; + +// Ordered log of the lifecycle the element drives on the inner: "reset" before each row, +// "infer" for each row. Mirrors the events-vector pattern in failsafe.cpp / accuracy_checked.cpp. +struct Recorder { + std::vector events; +}; + +class MockState : public ov::IVariableState { +public: + explicit MockState(std::shared_ptr rec) : ov::IVariableState("mock_state"), m_rec(std::move(rec)) {} + void reset() override { + m_rec->events.emplace_back("reset"); + } + void set_state(const ov::SoPtr&) override {} + ov::SoPtr get_state() const override { + return {}; + } + +private: + std::shared_ptr m_rec; +}; + +ov::Output port_named(const std::vector>& ports, const std::string& needle) { + for (const auto& port : ports) { + if (port.get_any_name().find(needle) != std::string::npos) { + return port; + } + } + OPENVINO_THROW("port not found: ", needle); +} + +std::string error_message(const ov::Exception& ex) { + return ex.what() == nullptr ? std::string{} : std::string(ex.what()); +} + +// Minimal two-output model (one input_ids input, two f32 outputs) for the multi-output test. The +// mock fills the outputs itself, so only the I/O signature matters. +std::shared_ptr build_two_output_model() { + auto ids = std::make_shared(ov::element::i64, ov::PartialShape{-1, -1}); + ids->set_friendly_name("input_ids"); + ids->output(0).set_names({"input_ids"}); + + auto as_f32 = std::make_shared(ids, ov::element::f32); + auto two = ov::op::v0::Constant::create(ov::element::f32, ov::Shape{}, {2.0f}); + auto scaled = std::make_shared(as_f32, two); + + return std::make_shared(ov::OutputVector{as_f32, scaled}, ov::ParameterVector{ids}, "two_output"); +} + +// Batch-1 inner stand-in: echoes the row's first input_ids token across its (generically +// shaped) output. Pure in its input, so the same row always yields the same value -- which +// is what lets callers check that batched scoring equals per-row scoring. +class MockInnerSync : public ov::ISyncInferRequest { +public: + MockInnerSync(std::shared_ptr compiled_model, std::shared_ptr rec) + : ov::ISyncInferRequest(std::move(compiled_model)), + m_rec(std::move(rec)), + m_state(std::make_shared(m_rec)) {} + + void infer() override { + m_rec->events.emplace_back("infer"); + + const auto ids = get_tensor(port_named(get_inputs(), "input_ids")); + const float token = static_cast(static_cast(ids->data())[0]); + + // Echo the row's token across every output; output k is offset by k so multi-output + // tests can tell the stacked outputs apart. + const auto& ports = get_outputs(); + for (std::size_t k = 0; k < ports.size(); ++k) { + const auto& port = ports[k]; + ov::Shape shape; + for (const auto& dim : port.get_partial_shape()) { + shape.push_back(dim.is_static() ? static_cast(dim.get_length()) : std::size_t{1}); + } + auto out = ov::get_tensor_impl(ov::Tensor(port.get_element_type(), shape)); + std::fill_n(static_cast(out->data()), + out->get_size(), + token + static_cast(k) * kOutputChannelOffset); + set_tensor(port, out); + } + } + + void check_tensors() const override {} + std::vector> query_state() const override { + return {ov::SoPtr(m_state)}; + } + std::vector get_profiling_info() const override { + return {}; + } + +private: + std::shared_ptr m_rec; + std::shared_ptr m_state; +}; + +class MockInnerCompiled : public ov::npuw::ICompiledModel { +public: + MockInnerCompiled(const std::shared_ptr& model, + const std::shared_ptr& plugin, + std::shared_ptr rec) + : ov::npuw::ICompiledModel(model, plugin), + m_rec(std::move(rec)) {} + + std::shared_ptr create_sync_infer_request() const override { + return std::make_shared(shared_from_this(), m_rec); + } + void export_model(std::ostream&) const override {} + std::shared_ptr get_runtime_model() const override { + return nullptr; + } + void set_property(const ov::AnyMap&) override {} + ov::Any get_property(const std::string&) const override { + return {}; + } + +private: + std::shared_ptr m_rec; +}; + +class NPUWBatchedElementTest : public ::testing::Test { +protected: + void SetUp() override { + m_plugin = std::make_shared(); + m_model = build_reranker_test_model(); + m_recorder = std::make_shared(); + } + + std::shared_ptr make_inner() { + return std::make_shared(m_model, m_plugin, m_recorder); + } + + std::shared_ptr wrap(bool enabled) { + return ov::npuw::batched::CompiledModel::create(make_inner(), m_plugin, enabled); + } + + // Bind a fresh zero-filled tensor of the port's element type to a named input. + template + static ov::SoPtr bind_input(const std::shared_ptr& req, + const std::string& name, + const ov::Shape& shape) { + const auto port = port_named(req->get_inputs(), name); + auto tensor = ov::get_tensor_impl(ov::Tensor(port.get_element_type(), shape)); + if (tensor->get_byte_size() > 0) { + std::memset(tensor->data(), 0, tensor->get_byte_size()); + } + req->set_tensor(port, tensor); + return tensor; + } + + // Set input_ids to [batch, len] and fill it row-major -- the only input the mock reads. + template + static void set_input_ids(const std::shared_ptr& req, const std::vector>& ids) { + const std::size_t batch = ids.size(); + const std::size_t len = ids.empty() ? 0 : ids.front().size(); + auto* tokens = static_cast(bind_input(req, "input_ids", {batch, len})->data()); + for (std::size_t r = 0; r < batch; ++r) { + std::copy(ids[r].begin(), ids[r].end(), tokens + r * len); + } + } + + // Bind the full reranker input set; the non-input_ids inputs just need a matching batch dim. + template + static void bind_inputs(const std::shared_ptr& req, const std::vector>& ids) { + const std::size_t batch = ids.size(); + const std::size_t len = ids.empty() ? 0 : ids.front().size(); + set_input_ids(req, ids); + bind_input(req, "attention_mask", {batch, len}); + bind_input(req, "position_ids", {batch, len}); + bind_input(req, "beam_idx", {batch}); + } + + // First element of output row r (the echoed token), independent of the output's rank. + static float row_value(const ov::SoPtr& out, std::size_t row) { + const std::size_t stride = out->get_size() / out->get_shape()[0]; + return static_cast(out->data())[row * stride]; + } + + std::shared_ptr m_plugin; + std::shared_ptr m_model; + std::shared_ptr m_recorder; +}; + +TEST_F(NPUWBatchedElementTest, RequestedFromProperties) { + EXPECT_FALSE(ov::npuw::batched::requested({})); + EXPECT_FALSE(ov::npuw::batched::requested({{"NPUW_TEXT_RERANK", false}})); + EXPECT_TRUE(ov::npuw::batched::requested({{"NPUW_TEXT_RERANK", true}})); + EXPECT_TRUE(ov::npuw::batched::requested({{"NPUW_TEXT_EMBED", true}})); + EXPECT_TRUE(ov::npuw::batched::requested({{"NPUW_LLM", true}, {"NPUW_TEXT_RERANK", true}})); +} + +TEST_F(NPUWBatchedElementTest, DisabledReturnsInnerUnwrapped) { + auto inner = make_inner(); + auto wrapped = ov::npuw::batched::CompiledModel::create(inner, m_plugin, /*enabled=*/false); + EXPECT_EQ(wrapped, inner); +} + +TEST_F(NPUWBatchedElementTest, EnabledWrapsInner) { + auto inner = make_inner(); + auto wrapped = ov::npuw::batched::CompiledModel::create(inner, m_plugin, /*enabled=*/true); + EXPECT_NE(wrapped, inner); + EXPECT_NE(std::dynamic_pointer_cast(wrapped), nullptr); + // The wrapper adds no I/O of its own -- it exposes the inner model's ports. + EXPECT_EQ(&wrapped->inputs(), &inner->inputs()); + EXPECT_EQ(&wrapped->outputs(), &inner->outputs()); +} + +TEST_F(NPUWBatchedElementTest, EachRowScoredIndependentlyAndStacked) { + auto wrapped = wrap(/*enabled=*/true); + auto req = wrapped->create_infer_request(); + + // Distinct first tokens so a misrouted row would surface as a wrong output value. + const std::vector> ids = {{11, 2, 3, 4}, {22, 5, 6, 7}, {33, 8, 9, 10}}; + bind_inputs(req, ids); + + req->infer(); + + const auto out = req->get_tensor(wrapped->outputs()[0]); + ASSERT_EQ(out->get_shape()[0], ids.size()); + for (std::size_t r = 0; r < ids.size(); ++r) { + EXPECT_FLOAT_EQ(row_value(out, r), static_cast(ids[r].front())) << "output row " << r; + } + + // One inner inference per row, each preceded by a state reset, in order. + EXPECT_EQ(m_recorder->events, (std::vector{"reset", "infer", "reset", "infer", "reset", "infer"})); +} + +TEST_F(NPUWBatchedElementTest, SingleRowScored) { + auto wrapped = wrap(/*enabled=*/true); + auto req = wrapped->create_infer_request(); + + const std::vector> ids = {{42, 8, 9, 10}}; + bind_inputs(req, ids); + + req->infer(); + + const auto out = req->get_tensor(wrapped->outputs()[0]); + ASSERT_EQ(out->get_shape()[0], 1u); + EXPECT_FLOAT_EQ(row_value(out, 0), static_cast(ids[0].front())); + EXPECT_EQ(m_recorder->events, (std::vector{"reset", "infer"})); +} + +// A zero-row batch has nothing to score and would publish unpopulated outputs. +TEST_F(NPUWBatchedElementTest, ZeroBatchIsRejected) { + auto wrapped = wrap(/*enabled=*/true); + auto req = wrapped->create_infer_request(); + + bind_inputs(req, {{1, 2, 3, 4}}); + bind_input(req, "input_ids", {0, 4}); + + try { + req->infer(); + FAIL() << "expected a zero-batch rejection"; + } catch (const ov::Exception& ex) { + EXPECT_NE(error_message(ex).find("zero-sized batch dimension"), std::string::npos); + } + EXPECT_TRUE(m_recorder->events.empty()); +} + +// input_ids fixes N=3; an attention_mask with batch 2 (neither N nor 1) cannot be sliced +// per row and must be rejected rather than fed whole into the batch-1 inner. +TEST_F(NPUWBatchedElementTest, MismatchedBatchRejected) { + auto wrapped = wrap(/*enabled=*/true); + auto req = wrapped->create_infer_request(); + + bind_inputs(req, {{1, 2}, {3, 4}, {5, 6}}); + bind_input(req, "attention_mask", {2, 2}); + + try { + req->infer(); + FAIL() << "expected a batch-mismatch rejection"; + } catch (const ov::Exception& ex) { + EXPECT_NE(error_message(ex).find("neither the inferred batch size"), std::string::npos); + } + EXPECT_TRUE(m_recorder->events.empty()); +} + +// A shared input (batch 1) is fed to every row whole -- only inputs carrying the full +// batch are sliced, so a [1, L] input is never sliced out of range. +TEST_F(NPUWBatchedElementTest, SharedInputBoundToEveryRow) { + auto wrapped = wrap(/*enabled=*/true); + auto req = wrapped->create_infer_request(); + + const std::vector> ids = {{11, 1, 1, 1}, {22, 1, 1, 1}, {33, 1, 1, 1}}; + bind_inputs(req, ids); + bind_input(req, "attention_mask", {1, 4}); // shared across the 3 rows + + req->infer(); + + const auto out = req->get_tensor(wrapped->outputs()[0]); + ASSERT_EQ(out->get_shape()[0], ids.size()); + for (std::size_t r = 0; r < ids.size(); ++r) { + EXPECT_FLOAT_EQ(row_value(out, r), static_cast(ids[r].front())) << "output row " << r; + } + EXPECT_EQ(m_recorder->events, (std::vector{"reset", "infer", "reset", "infer", "reset", "infer"})); +} + +// Regression: a leading batch-1 input (here input_ids, shared) must not pin the batch to 1 +// when a later input carries the real batch -- the batch is the largest leading dim. +TEST_F(NPUWBatchedElementTest, BatchSizeFromNonLeadingInput) { + auto wrapped = wrap(/*enabled=*/true); + auto req = wrapped->create_infer_request(); + + bind_inputs(req, {{7, 1, 1, 1}}); // input_ids = [1, 4], shared prompt (token 7) + bind_input(req, "attention_mask", {3, 4}); // real batch N = 3 on a non-leading input + + req->infer(); + + const auto out = req->get_tensor(wrapped->outputs()[0]); + ASSERT_EQ(out->get_shape()[0], 3u); + for (std::size_t r = 0; r < 3; ++r) { + EXPECT_FLOAT_EQ(row_value(out, r), 7.0f) << "output row " << r; // input_ids shared with every row + } + EXPECT_EQ(m_recorder->events, (std::vector{"reset", "infer", "reset", "infer", "reset", "infer"})); +} + +// Each model output is stacked into its own [N, ...] tensor. The mock offsets output k by k, so the +// two outputs must differ per row. +TEST_F(NPUWBatchedElementTest, MultipleOutputsStackedIndependently) { + auto model = build_two_output_model(); + auto inner = std::make_shared(model, m_plugin, m_recorder); + auto wrapped = ov::npuw::batched::CompiledModel::create(inner, m_plugin, /*enabled=*/true); + auto req = wrapped->create_infer_request(); + + const std::vector> ids = {{11, 1}, {22, 1}, {33, 1}}; + set_input_ids(req, ids); // the two-output model has only input_ids + + req->infer(); + + ASSERT_EQ(wrapped->outputs().size(), 2u); + const auto score = req->get_tensor(wrapped->outputs()[0]); + const auto hidden = req->get_tensor(wrapped->outputs()[1]); + ASSERT_EQ(score->get_shape()[0], ids.size()); + ASSERT_EQ(hidden->get_shape()[0], ids.size()); + for (std::size_t r = 0; r < ids.size(); ++r) { + EXPECT_FLOAT_EQ(row_value(score, r), static_cast(ids[r].front())); + EXPECT_FLOAT_EQ(row_value(hidden, r), static_cast(ids[r].front()) + kOutputChannelOffset); + } + EXPECT_EQ(m_recorder->events, (std::vector{"reset", "infer", "reset", "infer", "reset", "infer"})); +} + +// A caller-bound output tensor of the right shape and type is written in place, not replaced. +TEST_F(NPUWBatchedElementTest, PreBoundOutputTensorReused) { + auto model = build_two_output_model(); + auto inner = std::make_shared(model, m_plugin, m_recorder); + auto wrapped = ov::npuw::batched::CompiledModel::create(inner, m_plugin, /*enabled=*/true); + auto req = wrapped->create_infer_request(); + + const std::vector> ids = {{11, 1}, {22, 1}, {33, 1}}; + set_input_ids(req, ids); + + // The mock produces [1, 1] outputs (dynamic dims sized to 1), so the stacked shape is [3, 1]. + auto bound = ov::get_tensor_impl(ov::Tensor(ov::element::f32, {3, 1})); + req->set_tensor(wrapped->outputs()[0], bound); + + req->infer(); + + const auto out = req->get_tensor(wrapped->outputs()[0]); + EXPECT_EQ(out->data(), bound->data()); + for (std::size_t r = 0; r < ids.size(); ++r) { + EXPECT_FLOAT_EQ(row_value(out, r), static_cast(ids[r].front())); + } +} + +} // namespace diff --git a/src/plugins/intel_npu/tests/unit/npuw/llm_test_helpers.hpp b/src/plugins/intel_npu/tests/unit/npuw/llm_test_helpers.hpp index 48ce2997a6e1ec..a92672720351bb 100644 --- a/src/plugins/intel_npu/tests/unit/npuw/llm_test_helpers.hpp +++ b/src/plugins/intel_npu/tests/unit/npuw/llm_test_helpers.hpp @@ -69,8 +69,7 @@ inline std::shared_ptr build_dynamic_attention_llm_model() { for (const auto& input : model->inputs()) { const auto& name = input.get_any_name(); const auto& pshape = input.get_partial_shape(); - if (name.find("input_ids") != std::string::npos || - name.find("token_type_ids") != std::string::npos) { + if (name.find("input_ids") != std::string::npos || name.find("token_type_ids") != std::string::npos) { new_shapes[name] = ov::PartialShape{1, kSeq}; } else if (name.find("attention_mask") != std::string::npos) { new_shapes[name] = ov::PartialShape{1, kSeq + kPast}; @@ -100,6 +99,22 @@ inline std::shared_ptr build_llm_gqa_test_model() { return mb.build_llm(make_test_model_config_gqa()); } +/// Minimal Qwen3-style reranker: a GQA causal decoder with RMSNorm and per-head +/// Q/K normalization, stateful KV cache and an LM head (logits output). Matches the +/// I/O signature of Qwen3-Reranker (input_ids/attention_mask/position_ids + beam_idx), +/// which is what the batched scoring element fans out over. +inline LLMConfig make_test_model_config_reranker() { + auto cfg = make_test_model_config_gqa(); + cfg.norm = RMSNorm(cfg.hidden_size, cfg.precision); + cfg.qk_norm = RMSNorm(cfg.head_dim, cfg.precision); + return cfg; +} + +inline std::shared_ptr build_reranker_test_model() { + ModelBuilder mb; + return mb.build_llm(make_test_model_config_reranker()); +} + inline std::shared_ptr build_llm_test_model_with_kv_fake_convert(const ov::element::Type fake_convert_type) { auto model = build_llm_test_model(); auto scale = ov::op::v0::Constant::create(ov::element::f32, ov::Shape{}, {1.0f}); @@ -113,8 +128,7 @@ inline std::shared_ptr build_llm_test_model_with_kv_fake_convert(cons auto inject_fake_convert = [&](size_t input_idx, const std::string& suffix) { auto fake_convert_1 = std::make_shared(sdpa->input_value(input_idx), scale, fake_convert_type); - auto fake_convert_2 = - std::make_shared(fake_convert_1, scale, fake_convert_type); + auto fake_convert_2 = std::make_shared(fake_convert_1, scale, fake_convert_type); fake_convert_1->set_friendly_name(sdpa->get_friendly_name() + "/" + suffix + "_1"); fake_convert_2->set_friendly_name(sdpa->get_friendly_name() + "/" + suffix + "_2"); sdpa->input(input_idx).replace_source_output(fake_convert_2); @@ -328,8 +342,8 @@ class MockSubCompiledModel : public ov::npuw::ICompiledModel_v0 { }; struct CompileCall { - std::string friendly_name; - ov::AnyMap props; + std::string friendly_name; + ov::AnyMap props; std::shared_ptr model; };