From 740356318a2b2c003b013e72db85bf6979c2da7d Mon Sep 17 00:00:00 2001 From: Dylan Neve Date: Thu, 25 Jun 2026 16:21:49 +0100 Subject: [PATCH 1/7] NPUW: Add batched scoring as a composable v1 element Introduce ov::npuw::batched, a v1/elements decorator (alongside failsafe and accuracy_checked) that adds batch > 1 support on top of an inner compiled model that only supports a static batch size of 1, as NPUW's LLM and embedding pipelines do. 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 per-row outputs are stacked back into [N, ...] along axis 0. This is exact for single-shot scoring workloads with independent rows (text reranking, text embedding); generation is not a valid inner. infer() rejects a zero-sized batch and validates that every batched input's batch dimension is either N or 1 (shared/broadcast). The element composes on top of the other elements and the inner request needs no awareness of batching. create() returns the inner unwrapped when disabled. --- .../src/plugin/npuw/v1/elements/batched.cpp | 263 ++++++++++++++++++ .../src/plugin/npuw/v1/elements/batched.hpp | 124 +++++++++ 2 files changed, 387 insertions(+) create mode 100644 src/plugins/intel_npu/src/plugin/npuw/v1/elements/batched.cpp create mode 100644 src/plugins/intel_npu/src/plugin/npuw/v1/elements/batched.hpp 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 000000000000..6cad7aeed5ed --- /dev/null +++ b/src/plugins/intel_npu/src/plugin/npuw/v1/elements/batched.cpp @@ -0,0 +1,263 @@ +// Copyright (C) 2018-2026 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#include "batched.hpp" + +#include + +#include "openvino/core/except.hpp" +#include "openvino/runtime/make_tensor.hpp" +#include "openvino/runtime/tensor.hpp" + +namespace { + +// Zero-copy view of one batch row (axis 0) of a tensor: [N, ...] -> [1, ...]. +ov::SoPtr row_slice(const ov::SoPtr& tensor, std::size_t row) { + ov::Shape start_shape(tensor->get_shape().size(), 0u); + start_shape[0] = row; + ov::Shape end_shape = tensor->get_shape(); + end_shape[0] = row + 1; + return ov::get_tensor_impl(ov::Tensor(ov::make_tensor(tensor), start_shape, end_shape)); +} + +bool has_batch_dim(const ov::SoPtr& tensor) { + return !tensor->get_shape().empty(); +} + +} // namespace + +ov::SoPtr ov::npuw::batched::CompiledModel::create( + const std::shared_ptr& model, + const std::shared_ptr& plugin, + ov::SoPtr inner_compiled, + bool enabled) { + OPENVINO_ASSERT(inner_compiled._ptr != 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_compiled; + } + + auto compiled_model = std::make_shared(model, plugin, std::move(inner_compiled)); + return {compiled_model, {}}; +} + +ov::npuw::batched::CompiledModel::CompiledModel(const std::shared_ptr& model, + const std::shared_ptr& plugin, + ov::SoPtr inner_compiled) + : ov::ICompiledModel(model, plugin), + m_inner(std::move(inner_compiled)) { + OPENVINO_ASSERT(m_inner._ptr != nullptr, "Batched compiled model requires an inner compiled model"); +} + +void ov::npuw::batched::CompiledModel::export_model(std::ostream& model) const { + 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); +} + +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)); +} + +std::shared_ptr ov::npuw::batched::CompiledModel::create_infer_request() const { + return std::make_shared(create_sync_infer_request(), + get_task_executor(), + get_callback_executor()); +} + +ov::npuw::batched::InferRequest::InferRequest(const std::shared_ptr& compiled_model, + std::shared_ptr inner_request) + : ov::ISyncInferRequest(compiled_model), + m_inner_async(std::move(inner_request)) { + OPENVINO_ASSERT(m_inner_async != nullptr, "Batched element requires a non-null inner request"); + init_public_tensors(); +} + +ov::npuw::batched::InferRequest::InferRequest(const std::shared_ptr& compiled_model, + std::shared_ptr inner_request) + : ov::ISyncInferRequest(compiled_model), + m_inner_sync(std::move(inner_request)) { + OPENVINO_ASSERT(m_inner_sync != nullptr, "Batched element requires a non-null inner request"); + init_public_tensors(); +} + +void ov::npuw::batched::InferRequest::init_public_tensors() { + // Allocate a tensor for every public port so get_tensor() never returns an + // uninitialized handle (callers such as the Python infer(dict) dispatcher fetch + // the tensor before populating it). Dynamic dims are sized to 0 and resized later; + // the real [N, ...] tensors are bound by the caller (inputs) or by infer() (outputs). + const auto init_port = [this](const ov::Output& port) { + if (ov::ISyncInferRequest::get_tensor(port)) { + return; + } + const auto& pshape = port.get_partial_shape(); + ov::Shape shape; + if (pshape.is_dynamic()) { + for (const auto& dim : pshape) { + shape.push_back(dim.is_static() ? dim.get_length() : 0); + } + } else { + shape = pshape.to_shape(); + } + set_tensor(port, ov::get_tensor_impl(ov::Tensor(port.get_element_type(), shape))); + }; + for (const auto& port : get_inputs()) { + init_port(port); + } + for (const auto& port : get_outputs()) { + init_port(port); + } +} + +const std::vector>& ov::npuw::batched::InferRequest::inner_inputs() const { + return m_inner_sync ? m_inner_sync->get_inputs() : m_inner_async->get_inputs(); +} + +const std::vector>& ov::npuw::batched::InferRequest::inner_outputs() const { + return m_inner_sync ? m_inner_sync->get_outputs() : m_inner_async->get_outputs(); +} + +void ov::npuw::batched::InferRequest::inner_set_tensor(const ov::Output& port, + const ov::SoPtr& tensor) { + if (m_inner_sync) { + m_inner_sync->set_tensor(port, tensor); + } else { + m_inner_async->set_tensor(port, tensor); + } +} + +ov::SoPtr ov::npuw::batched::InferRequest::inner_get_tensor(const ov::Output& port) const { + return m_inner_sync ? m_inner_sync->get_tensor(port) : m_inner_async->get_tensor(port); +} + +void ov::npuw::batched::InferRequest::inner_infer() { + if (m_inner_sync) { + m_inner_sync->infer(); + } else { + m_inner_async->infer(); + } +} + +std::vector> ov::npuw::batched::InferRequest::inner_query_state() const { + return m_inner_sync ? m_inner_sync->query_state() : m_inner_async->query_state(); +} + +void ov::npuw::batched::InferRequest::infer() { + std::lock_guard lock(m_mutex); + + const auto& wrapper_inputs = get_inputs(); + const auto& wrapper_outputs = get_outputs(); + const auto& in_ports = inner_inputs(); + const auto& out_ports = inner_outputs(); + + OPENVINO_ASSERT(wrapper_inputs.size() == in_ports.size() && wrapper_outputs.size() == out_ports.size(), + "Batched element: inner request I/O does not match the wrapped model"); + + // Batch size is taken from the first input that carries a batch dimension. + std::size_t batch = 1; + for (const auto& port : wrapper_inputs) { + const auto tensor = get_tensor(port); + if (has_batch_dim(tensor)) { + batch = tensor->get_shape()[0]; + break; + } + } + // A zero-sized batch has no rows to score and would leave the outputs unpopulated + // (publishing null tensors). Reject it explicitly rather than silently no-op. + OPENVINO_ASSERT(batch > 0, "Batched element: batch size must be > 0, got an input with batch dimension 0."); + + // Validate every batched input up front: each row-carrying input must have a batch + // dimension equal to `batch`, or exactly 1 (a shared/broadcast input passed to every row). + // Anything else (e.g. a [M, ...] input with M != batch and M != 1) cannot be sliced per row + // and would otherwise be fed whole into the batch-1 inner request -> wrong results. + for (std::size_t i = 0; i < wrapper_inputs.size(); ++i) { + const auto full = get_tensor(wrapper_inputs[i]); + if (!has_batch_dim(full)) { + continue; + } + const std::size_t in_batch = full->get_shape()[0]; + OPENVINO_ASSERT(in_batch == batch || in_batch == 1, + "Batched element: input '", + wrapper_inputs[i].get_any_name(), + "' has batch dimension ", + in_batch, + " which is neither the inferred batch size ", + batch, + " nor 1 (broadcast)."); + } + + // Aggregated [batch, ...] outputs, allocated lazily once the per-row output shape is known. + std::vector> aggregated_outputs(wrapper_outputs.size()); + + // The inner request's variable states are stable across rows, so query them once and reset + // them per row rather than re-querying (which may allocate) on every iteration. + const auto inner_states = inner_query_state(); + + for (std::size_t row = 0; row < batch; ++row) { + // Rows are independent prompts: clear the inner request's variable state + // (KV-cache) so row i never sees row i-1. Harmless for stateless inners. + for (const auto& state : inner_states) { + state->reset(); + } + + // Bind the row's slice of every batched input; pass shared ([1, ...] or non-batched) + // inputs through unchanged. + for (std::size_t i = 0; i < wrapper_inputs.size(); ++i) { + const auto full = get_tensor(wrapper_inputs[i]); + const bool sliceable = batch > 1 && has_batch_dim(full) && full->get_shape()[0] == batch; + inner_set_tensor(in_ports[i], sliceable ? row_slice(full, row) : full); + } + + inner_infer(); + + // Stack each per-row output into row i of the aggregated [batch, ...] tensor. + for (std::size_t i = 0; i < wrapper_outputs.size(); ++i) { + const auto inner_out = inner_get_tensor(out_ports[i]); + if (!aggregated_outputs[i]) { + ov::Shape out_shape = inner_out->get_shape(); + if (!out_shape.empty()) { + out_shape[0] = batch; + } + aggregated_outputs[i] = ov::get_tensor_impl(ov::Tensor(inner_out->get_element_type(), out_shape)); + } + if (batch > 1 && has_batch_dim(inner_out)) { + const auto slot = row_slice(aggregated_outputs[i], row); + inner_out->copy_to(slot._ptr); + } else { + inner_out->copy_to(aggregated_outputs[i]._ptr); + } + } + } + + // Publish the aggregated outputs as this request's public output tensors. + for (std::size_t i = 0; i < wrapper_outputs.size(); ++i) { + set_tensor(wrapper_outputs[i], aggregated_outputs[i]); + } +} + +void ov::npuw::batched::InferRequest::check_tensors() const {} + +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 {}; +} 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 000000000000..d73455737606 --- /dev/null +++ b/src/plugins/intel_npu/src/plugin/npuw/v1/elements/batched.hpp @@ -0,0 +1,124 @@ +// Copyright (C) 2018-2026 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#pragma once + +#include +#include +#include + +#include "openvino/runtime/iasync_infer_request.hpp" +#include "openvino/runtime/icompiled_model.hpp" +#include "openvino/runtime/isync_infer_request.hpp" +#include "openvino/runtime/so_ptr.hpp" + +namespace ov::npuw::batched { + +class InferRequest; + +// A compiled-model wrapper 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 ov::ICompiledModel decorator that exposes the same I/O as the +// model it wraps and forwards property queries 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 stacked +// back into [N, ...] tensors along the batch dimension (axis 0). +// +// 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 rows are not independent. +// +// It composes on top of the other elements, e.g.: +// +// Batched( AccuracyChecked( Failsafe(NPU -> CPU) ) ) +// +// and the inner request needs no awareness of batching at all -- a stock +// single-sequence LLM/embedding infer request works unchanged. +class CompiledModel final : public ov::ICompiledModel { +public: + // Factory method. Returns inner_compiled unwrapped when enabled == false, + // keeping the zero-overhead path trivial (mirrors accuracy_checked::create). + static ov::SoPtr create(const std::shared_ptr& model, + const std::shared_ptr& plugin, + ov::SoPtr inner_compiled, + bool enabled); + + CompiledModel(const std::shared_ptr& model, + const std::shared_ptr& plugin, + ov::SoPtr inner_compiled); + + 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; + + std::shared_ptr create_sync_infer_request() const override; + std::shared_ptr create_infer_request() const override; + +private: + friend class InferRequest; + + ov::SoPtr m_inner; +}; + +// Sync infer request that unrolls a batched inference over a single-sequence +// inner request. +// +// On each infer() call it determines the batch size N from the inputs, then for +// each row 0..N-1: resets the inner request's variable state, binds the row's +// [1, ...] slice of every input, runs the inner request, and copies the inner +// [1, ...] output into row i of the wrapper's [N, ...] output tensors. The +// public input/output tensors are held by the ISyncInferRequest base; only the +// inner request sees [1, ...] shapes. +// +// It is constructed from the compiled model whose I/O it exposes plus the inner +// request to drive. This lets it be produced both by Batched::CompiledModel +// (standalone, composable element) and reused directly by a pipeline that +// already owns a single-sequence request (e.g. NPUW's LLMCompiledModel wrapping +// its LLMInferRequest for scoring), without that pipeline having to wrap its +// whole compiled model. +class InferRequest final : public ov::ISyncInferRequest { +public: + // Drive an async inner request (used by the standalone Batched::CompiledModel decorator, + // where the inner is a separate compiled model with its own task executor). + InferRequest(const std::shared_ptr& compiled_model, + std::shared_ptr inner_request); + + // Drive a sync inner request directly, on the calling thread, without an executor. Used when a + // pipeline wraps its own single-sequence request (e.g. NPUW's LLMCompiledModel wrapping its + // LLMInferRequest): this avoids a nested-executor deadlock that would otherwise occur if the + // inner ran on the same task executor as the outer (async) request. + 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: + void init_public_tensors(); + + // The inner single-sequence request, driven row by row. Exactly one of the two handles is set; + // the small accessors below hide which interface (sync/async) is in use from infer(). + const std::vector>& inner_inputs() const; + const std::vector>& inner_outputs() const; + void inner_set_tensor(const ov::Output& port, const ov::SoPtr& tensor); + ov::SoPtr inner_get_tensor(const ov::Output& port) const; + void inner_infer(); + std::vector> inner_query_state() const; + + std::shared_ptr m_inner_async; + std::shared_ptr m_inner_sync; + mutable std::mutex m_mutex; +}; + +} // namespace ov::npuw::batched From a88878274d8e127b06490932a1dfbf52da529d45 Mon Sep 17 00:00:00 2001 From: Dylan Neve Date: Thu, 25 Jun 2026 16:21:49 +0100 Subject: [PATCH 2/7] NPUW: Wire batched scoring into LLMCompiledModel, gated on model kind Wrap the single-sequence request created by LLMCompiledModel with the batched element so a [N, ...] scoring inference is unrolled into N independent [1, ...] inferences and stacked back. This makes GenAI's text reranking and text embedding pipelines work for more than one document/text on NPU through the unchanged batch-1 compiled graphs, with no application-side per-row loop. The element is enabled automatically for the non-generating scoring pipelines -- text embedding (m_is_embedding, already known) and text reranking (m_is_rerank, set from a new NPUW_TEXT_RERANK key mirroring NPUW_TEXT_EMBED) -- rather than via an application-facing option. A reranker is a plain causal LM, indistinguishable from chat at the plugin level, so the pipeline tags it; the KV-reset-between-rows element is therefore never wired into generation. The inner request is driven synchronously so it runs on the calling thread instead of being re-scheduled onto this model's task executor (nesting deadlocks). --- .../intel_npu/config/npuw_option_defs.inc | 1 + .../src/plugin/npuw/llm_compiled_model.cpp | 32 ++++++++++++++++--- .../src/plugin/npuw/llm_compiled_model.hpp | 1 + 3 files changed, 29 insertions(+), 5 deletions(-) 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 081be721abc4..81e187efd1d1 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/llm_compiled_model.cpp b/src/plugins/intel_npu/src/plugin/npuw/llm_compiled_model.cpp index 53b9059d03f7..8473a7ee27e2 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 @@ -10,6 +10,7 @@ #include "llm_compiled_model_utils.hpp" #include "llm_infer_request.hpp" #include "logging.hpp" +#include "v1/elements/batched.hpp" #include "moe_transformations/apply_moe_device_routed_transforms.hpp" #include "npuw_transformations/add_position_ids_param.hpp" #include "npuw_transformations/convert_kvcache_to_precision.hpp" @@ -815,6 +816,9 @@ 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; + auto use_text_rerank_key = pop_option(other_props, std::string("NPUW_TEXT_RERANK")); + m_is_rerank = use_text_rerank_key.value_or(false).as() == true; + if (m_is_embedding) { LOG_DEBUG("Text-embedding model rebuild"); ov::npuw::util::PrepareTextEmbeddingModel(seq_len_dim).run_on_model(kvcache_model); @@ -1619,13 +1623,30 @@ ov::Any ov::npuw::LLMCompiledModel::get_property(const std::string& name) const std::shared_ptr ov::npuw::LLMCompiledModel::create_sync_infer_request() const { auto* non_const_this = const_cast(this); // because of const in API + if (m_is_whisper) { return non_const_this->create_whisper_infer_request(); - } else if (m_is_embedding) { - return non_const_this->create_embedding_infer_request(); - } else { - return non_const_this->create_llm_infer_request(); } + + auto inner = m_is_embedding ? non_const_this->create_embedding_infer_request() + : non_const_this->create_llm_infer_request(); + + // Batched scoring: wrap the single-sequence request with the batched element so a + // [N, ...] input is unrolled into N independent [1, ...] inferences (rows are scored + // one at a time, with the KV-cache reset between them) and stacked back into [N, ...]. + // This is valid only for the non-generating scoring pipelines -- text embedding and + // text reranking -- where rows are independent; the pipeline flags the model as such at + // compile time, so no user-facing option is needed. Generation never sets these. + if (m_is_embedding || m_is_rerank) { + // Drive the single-sequence request synchronously (sync overload) so it runs on the + // calling thread rather than being re-scheduled onto this model's task executor -- the + // outer request is already async on that executor, and nesting on the same executor + // would deadlock. + auto self = std::static_pointer_cast(shared_from_this()); + return std::make_shared(self, std::move(inner)); + } + + return inner; } std::shared_ptr ov::npuw::LLMCompiledModel::create_llm_infer_request() { @@ -1675,6 +1696,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/llm_compiled_model.hpp b/src/plugins/intel_npu/src/plugin/npuw/llm_compiled_model.hpp index d9796d2d444a..74fd777eedad 100644 --- a/src/plugins/intel_npu/src/plugin/npuw/llm_compiled_model.hpp +++ b/src/plugins/intel_npu/src/plugin/npuw/llm_compiled_model.hpp @@ -142,6 +142,7 @@ class LLMCompiledModel : public ov::npuw::ICompiledModel { size_t m_decomposed_sdpa_size = 0; bool m_is_embedding = false; + bool m_is_rerank = false; // Create generate model variants with different sizes std::vector> create_generate_model_variants( From 80c205ebbf0d300921d73f03b6bff4b8a1abadd9 Mon Sep 17 00:00:00 2001 From: Dylan Neve Date: Thu, 25 Jun 2026 16:21:49 +0100 Subject: [PATCH 3/7] NPUW: Unit tests for the batched element Add unit tests for ov::npuw::batched driving the element through a synthetic Qwen-style reranker (build_reranker_test_model, added to llm_test_helpers.hpp: GQA + RMSNorm + per-head QK-norm) and a recording mock inner. The element is the unit under test; only the inner compiled model is mocked, keeping the tests deviceless and able to observe the per-row infer/reset. The mock echoes each row's first token, so the tests assert the element's core guarantee directly -- output row r reflects input row r -- across the disabled, enabled, multi-row, single-row, zero-batch and batch-mismatch paths. Following failsafe.cpp / accuracy_checked.cpp, an ordered event log checks the per-row state reset and the rejection paths assert the error message. --- .../intel_npu/tests/unit/CMakeLists.txt | 2 + .../tests/unit/npuw/batched_element_test.cpp | 275 ++++++++++++++++++ .../tests/unit/npuw/llm_test_helpers.hpp | 16 + 3 files changed, 293 insertions(+) create mode 100644 src/plugins/intel_npu/tests/unit/npuw/batched_element_test.cpp diff --git a/src/plugins/intel_npu/tests/unit/CMakeLists.txt b/src/plugins/intel_npu/tests/unit/CMakeLists.txt index 9dc2cf8e2f00..123e8315db11 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 000000000000..7a9bb9e5c8b4 --- /dev/null +++ b/src/plugins/intel_npu/tests/unit/npuw/batched_element_test.cpp @@ -0,0 +1,275 @@ +// 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 stacks the per-row outputs back to [N, ...]. +// 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 "llm_test_helpers.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; + +// 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()); +} + +// 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]); + + const auto port = get_outputs()[0]; + 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); + 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::ICompiledModel { +public: + MockInnerCompiled(const std::shared_ptr& model, + const std::shared_ptr& plugin, + std::shared_ptr rec) + : ov::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); + } + std::shared_ptr create_infer_request() const override { + return std::make_shared(create_sync_infer_request(), + get_task_executor(), + get_callback_executor()); + } + 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(); + } + + ov::SoPtr make_inner() { + return {std::make_shared(m_model, m_plugin, m_recorder), {}}; + } + + ov::SoPtr wrap(bool enabled) { + return ov::npuw::batched::CompiledModel::create(m_model, m_plugin, make_inner(), enabled); + } + + // Resize a named input on the request, in place, and return it. + static ov::SoPtr resize_input(const std::shared_ptr& req, + const ov::SoPtr& wrapped, + const std::string& name, + const ov::Shape& shape) { + const auto tensor = req->get_tensor(port_named(wrapped->inputs(), name)); + tensor->set_shape(shape); + return tensor; + } + + // Bind the reranker inputs at the batch/length implied by `ids`. Only input_ids carries + // values the mock reads; attention_mask, position_ids and beam_idx just need a matching + // batch dimension, so they are sized but left unset. + static void bind_inputs(const std::shared_ptr& req, + const ov::SoPtr& wrapped, + 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(resize_input(req, wrapped, "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); + } + resize_input(req, wrapped, "attention_mask", {batch, len}); + resize_input(req, wrapped, "position_ids", {batch, len}); + resize_input(req, wrapped, "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, DisabledReturnsInnerUnwrapped) { + auto inner = make_inner(); + auto wrapped = ov::npuw::batched::CompiledModel::create(m_model, m_plugin, inner, /*enabled=*/false); + EXPECT_EQ(wrapped._ptr, inner._ptr); +} + +TEST_F(NPUWBatchedElementTest, EnabledWrapsInner) { + auto inner = make_inner(); + auto wrapped = ov::npuw::batched::CompiledModel::create(m_model, m_plugin, inner, /*enabled=*/true); + EXPECT_NE(wrapped._ptr, inner._ptr); + EXPECT_NE(std::dynamic_pointer_cast(wrapped._ptr), nullptr); +} + +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, wrapped, 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, wrapped, 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 an unpopulated output. +TEST_F(NPUWBatchedElementTest, ZeroBatchIsRejected) { + auto wrapped = wrap(/*enabled=*/true); + auto req = wrapped->create_infer_request(); + + resize_input(req, wrapped, "input_ids", {0, 4}); + + try { + req->infer(); + FAIL() << "expected a zero-batch rejection"; + } catch (const ov::Exception& ex) { + EXPECT_NE(error_message(ex).find("batch size must be > 0"), 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, wrapped, {{1, 2}, {3, 4}, {5, 6}}); + resize_input(req, wrapped, "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()); +} + +} // 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 48ce2997a6e1..2f638e024cf9 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 @@ -100,6 +100,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}); From 35da727f66c34404ec23511fdf0e0bc597d5c10b Mon Sep 17 00:00:00 2001 From: Dylan Neve Date: Fri, 26 Jun 2026 10:55:13 +0100 Subject: [PATCH 4/7] NPUW: Batched element review fixes -- broadcast-aware detection, single-row fast path, more tests Address review feedback on the batched scoring element: - Detect the batch as the largest leading dimension across inputs instead of the first input that has one, so a leading [1, ...] broadcast input no longer pins the batch to 1 when a non-leading input carries the real batch N. - Skip the aggregation buffer and per-row copy for a single row (publish the inner output as-is). - Snapshot the public input tensors once per infer() rather than re-fetching them for detection, validation, and per-row binding. - Tighten the comments on the empty check_tensors() and the output stacking. Tests: - SyncInnerConstructorScoresEachRow: exercises the sync-inner constructor used by the production LLM/rerank wiring (previously only the async-inner path was covered). - BroadcastInputSharedAcrossRows: a [1, ...] input is fed to every row unsliced. - BatchSizeFromNonLeadingInput: regression for the detection fix above. - MultipleOutputsStackedIndependently: each model output is stacked into its own [N, ...] tensor. All NPUWBatchedElementTest cases pass. --- .../src/plugin/npuw/v1/elements/batched.cpp | 78 ++++---- .../src/plugin/npuw/v1/elements/batched.hpp | 13 +- .../tests/unit/npuw/batched_element_test.cpp | 179 +++++++++++++++--- 3 files changed, 204 insertions(+), 66 deletions(-) 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 index 6cad7aeed5ed..1bd8f7f80391 100644 --- a/src/plugins/intel_npu/src/plugin/npuw/v1/elements/batched.cpp +++ b/src/plugins/intel_npu/src/plugin/npuw/v1/elements/batched.cpp @@ -4,6 +4,7 @@ #include "batched.hpp" +#include #include #include "openvino/core/except.hpp" @@ -168,29 +169,35 @@ void ov::npuw::batched::InferRequest::infer() { OPENVINO_ASSERT(wrapper_inputs.size() == in_ports.size() && wrapper_outputs.size() == out_ports.size(), "Batched element: inner request I/O does not match the wrapped model"); - // Batch size is taken from the first input that carries a batch dimension. - std::size_t batch = 1; + // Snapshot the public inputs once -- read repeatedly below. + std::vector> in_tensors; + in_tensors.reserve(wrapper_inputs.size()); for (const auto& port : wrapper_inputs) { - const auto tensor = get_tensor(port); - if (has_batch_dim(tensor)) { - batch = tensor->get_shape()[0]; - break; + in_tensors.push_back(get_tensor(port)); + } + + // Batch is the largest leading dim across inputs, so a shared [1, ...] input is broadcast + // rather than mistaken for the batch. Stays 1 when nothing is batched. + std::size_t batch = 1; + bool any_batched = false; + for (const auto& tensor : in_tensors) { + if (!has_batch_dim(tensor)) { + continue; } + const std::size_t in_batch = tensor->get_shape()[0]; + batch = any_batched ? std::max(batch, in_batch) : in_batch; + any_batched = true; } - // A zero-sized batch has no rows to score and would leave the outputs unpopulated - // (publishing null tensors). Reject it explicitly rather than silently no-op. + // A zero-sized batch has no rows and would publish unpopulated outputs; reject it. OPENVINO_ASSERT(batch > 0, "Batched element: batch size must be > 0, got an input with batch dimension 0."); - // Validate every batched input up front: each row-carrying input must have a batch - // dimension equal to `batch`, or exactly 1 (a shared/broadcast input passed to every row). - // Anything else (e.g. a [M, ...] input with M != batch and M != 1) cannot be sliced per row - // and would otherwise be fed whole into the batch-1 inner request -> wrong results. + // Every batched input must match `batch` or be a broadcast (dim 0 == 1); anything else + // cannot be sliced per row. for (std::size_t i = 0; i < wrapper_inputs.size(); ++i) { - const auto full = get_tensor(wrapper_inputs[i]); - if (!has_batch_dim(full)) { + if (!has_batch_dim(in_tensors[i])) { continue; } - const std::size_t in_batch = full->get_shape()[0]; + const std::size_t in_batch = in_tensors[i]->get_shape()[0]; OPENVINO_ASSERT(in_batch == batch || in_batch == 1, "Batched element: input '", wrapper_inputs[i].get_any_name(), @@ -201,33 +208,38 @@ void ov::npuw::batched::InferRequest::infer() { " nor 1 (broadcast)."); } - // Aggregated [batch, ...] outputs, allocated lazily once the per-row output shape is known. - std::vector> aggregated_outputs(wrapper_outputs.size()); - - // The inner request's variable states are stable across rows, so query them once and reset - // them per row rather than re-querying (which may allocate) on every iteration. + // States are stable across rows: query once, reset per row so row i never sees row i-1. const auto inner_states = inner_query_state(); - - for (std::size_t row = 0; row < batch; ++row) { - // Rows are independent prompts: clear the inner request's variable state - // (KV-cache) so row i never sees row i-1. Harmless for stateless inners. + const auto reset_inner_state = [&] { for (const auto& state : inner_states) { state->reset(); } + }; - // Bind the row's slice of every batched input; pass shared ([1, ...] or non-batched) - // inputs through unchanged. + // Slice per-row inputs out of [batch, ...]; pass broadcast/non-batched inputs through whole. + const auto bind_row = [&](std::size_t row) { for (std::size_t i = 0; i < wrapper_inputs.size(); ++i) { - const auto full = get_tensor(wrapper_inputs[i]); + const auto& full = in_tensors[i]; const bool sliceable = batch > 1 && has_batch_dim(full) && full->get_shape()[0] == batch; inner_set_tensor(in_ports[i], sliceable ? row_slice(full, row) : full); } + }; + + // Per-row outputs stacked along axis 0. A single row has nothing to stack, so its output is + // published as-is (no aggregation buffer or copy); a non-batched output collapses to the last row. + std::vector> aggregated_outputs(wrapper_outputs.size()); + for (std::size_t row = 0; row < batch; ++row) { + reset_inner_state(); + bind_row(row); inner_infer(); - // Stack each per-row output into row i of the aggregated [batch, ...] tensor. for (std::size_t i = 0; i < wrapper_outputs.size(); ++i) { const auto inner_out = inner_get_tensor(out_ports[i]); + if (batch == 1) { + aggregated_outputs[i] = inner_out; + continue; + } if (!aggregated_outputs[i]) { ov::Shape out_shape = inner_out->get_shape(); if (!out_shape.empty()) { @@ -235,22 +247,22 @@ void ov::npuw::batched::InferRequest::infer() { } aggregated_outputs[i] = ov::get_tensor_impl(ov::Tensor(inner_out->get_element_type(), out_shape)); } - if (batch > 1 && has_batch_dim(inner_out)) { - const auto slot = row_slice(aggregated_outputs[i], row); - inner_out->copy_to(slot._ptr); + if (has_batch_dim(inner_out)) { + inner_out->copy_to(row_slice(aggregated_outputs[i], row)._ptr); } else { inner_out->copy_to(aggregated_outputs[i]._ptr); } } } - // Publish the aggregated outputs as this request's public output tensors. for (std::size_t i = 0; i < wrapper_outputs.size(); ++i) { set_tensor(wrapper_outputs[i], aggregated_outputs[i]); } } -void ov::npuw::batched::InferRequest::check_tensors() const {} +void ov::npuw::batched::InferRequest::check_tensors() const { + // No-op: batched tensors are unrolled and validated per row by the inner request. +} std::vector> ov::npuw::batched::InferRequest::query_state() const { // The batched element resets inner state between rows and exposes no 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 index d73455737606..204f551ce0f8 100644 --- a/src/plugins/intel_npu/src/plugin/npuw/v1/elements/batched.hpp +++ b/src/plugins/intel_npu/src/plugin/npuw/v1/elements/batched.hpp @@ -71,12 +71,13 @@ class CompiledModel final : public ov::ICompiledModel { // Sync infer request that unrolls a batched inference over a single-sequence // inner request. // -// On each infer() call it determines the batch size N from the inputs, then for -// each row 0..N-1: resets the inner request's variable state, binds the row's -// [1, ...] slice of every input, runs the inner request, and copies the inner -// [1, ...] output into row i of the wrapper's [N, ...] output tensors. The -// public input/output tensors are held by the ISyncInferRequest base; only the -// inner request sees [1, ...] shapes. +// On each infer() call it takes the batch size N as the largest leading dimension +// across the inputs (so a shared [1, ...] input is broadcast, not mistaken for the +// batch), then for each row 0..N-1: resets the inner request's variable state, binds +// the row's [1, ...] slice of every per-row input, runs the inner request, and copies +// the inner [1, ...] output into row i of the wrapper's [N, ...] output tensors. +// N == 1 publishes the inner outputs directly. The public input/output tensors are +// held by the ISyncInferRequest base; only the inner request sees [1, ...] shapes. // // It is constructed from the compiled model whose I/O it exposes plus the inner // request to drive. This lets it be produced both by Batched::CompiledModel 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 index 7a9bb9e5c8b4..2b7d38e39794 100644 --- a/src/plugins/intel_npu/tests/unit/npuw/batched_element_test.cpp +++ b/src/plugins/intel_npu/tests/unit/npuw/batched_element_test.cpp @@ -19,6 +19,11 @@ #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" @@ -31,6 +36,9 @@ 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 { @@ -66,6 +74,20 @@ 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. @@ -82,14 +104,21 @@ class MockInnerSync : public ov::ISyncInferRequest { const auto ids = get_tensor(port_named(get_inputs(), "input_ids")); const float token = static_cast(static_cast(ids->data())[0]); - const auto port = get_outputs()[0]; - 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}); + // 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); } - auto out = ov::get_tensor_impl(ov::Tensor(port.get_element_type(), shape)); - std::fill_n(static_cast(out->data()), out->get_size(), token); - set_tensor(port, out); } void check_tensors() const override {} @@ -150,32 +179,37 @@ class NPUWBatchedElementTest : public ::testing::Test { return ov::npuw::batched::CompiledModel::create(m_model, m_plugin, make_inner(), enabled); } - // Resize a named input on the request, in place, and return it. - static ov::SoPtr resize_input(const std::shared_ptr& req, - const ov::SoPtr& wrapped, + // Resize a named input in place and return it. Templated so it serves both the async wrapper + // and a bare sync request -- both expose get_inputs()/get_tensor(). + template + static ov::SoPtr resize_input(const std::shared_ptr& req, const std::string& name, const ov::Shape& shape) { - const auto tensor = req->get_tensor(port_named(wrapped->inputs(), name)); + const auto tensor = req->get_tensor(port_named(req->get_inputs(), name)); tensor->set_shape(shape); return tensor; } - // Bind the reranker inputs at the batch/length implied by `ids`. Only input_ids carries - // values the mock reads; attention_mask, position_ids and beam_idx just need a matching - // batch dimension, so they are sized but left unset. - static void bind_inputs(const std::shared_ptr& req, - const ov::SoPtr& wrapped, - const std::vector>& ids) { + // 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(resize_input(req, wrapped, "input_ids", {batch, len})->data()); + auto* tokens = static_cast(resize_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); } - resize_input(req, wrapped, "attention_mask", {batch, len}); - resize_input(req, wrapped, "position_ids", {batch, len}); - resize_input(req, wrapped, "beam_idx", {batch}); + } + + // 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); + resize_input(req, "attention_mask", {batch, len}); + resize_input(req, "position_ids", {batch, len}); + resize_input(req, "beam_idx", {batch}); } // First element of output row r (the echoed token), independent of the output's rank. @@ -208,7 +242,7 @@ TEST_F(NPUWBatchedElementTest, EachRowScoredIndependentlyAndStacked) { // 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, wrapped, ids); + bind_inputs(req, ids); req->infer(); @@ -228,7 +262,7 @@ TEST_F(NPUWBatchedElementTest, SingleRowScored) { auto req = wrapped->create_infer_request(); const std::vector> ids = {{42, 8, 9, 10}}; - bind_inputs(req, wrapped, ids); + bind_inputs(req, ids); req->infer(); @@ -243,7 +277,7 @@ TEST_F(NPUWBatchedElementTest, ZeroBatchIsRejected) { auto wrapped = wrap(/*enabled=*/true); auto req = wrapped->create_infer_request(); - resize_input(req, wrapped, "input_ids", {0, 4}); + resize_input(req, "input_ids", {0, 4}); try { req->infer(); @@ -260,8 +294,8 @@ TEST_F(NPUWBatchedElementTest, MismatchedBatchRejected) { auto wrapped = wrap(/*enabled=*/true); auto req = wrapped->create_infer_request(); - bind_inputs(req, wrapped, {{1, 2}, {3, 4}, {5, 6}}); - resize_input(req, wrapped, "attention_mask", {2, 2}); + bind_inputs(req, {{1, 2}, {3, 4}, {5, 6}}); + resize_input(req, "attention_mask", {2, 2}); try { req->infer(); @@ -272,4 +306,95 @@ TEST_F(NPUWBatchedElementTest, MismatchedBatchRejected) { EXPECT_TRUE(m_recorder->events.empty()); } +// The production LLM/rerank wiring builds the element from a *sync* inner request (driven on the +// calling thread to dodge a nested-executor deadlock), bypassing batched::CompiledModel. Exercise +// that constructor directly. +TEST_F(NPUWBatchedElementTest, SyncInnerConstructorScoresEachRow) { + auto inner_compiled = std::make_shared(m_model, m_plugin, m_recorder); + std::shared_ptr inner_sync = inner_compiled->create_sync_infer_request(); + std::shared_ptr compiled = inner_compiled; + auto req = std::make_shared(compiled, std::move(inner_sync)); + + 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(req->get_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"})); +} + +// A shared/broadcast input (batch 1) is fed to every row unsliced -- the slice guard only fires +// when the leading dim equals the batch, so a [1, L] input is never sliced out of range. +TEST_F(NPUWBatchedElementTest, BroadcastInputSharedAcrossRows) { + 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); + resize_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 old "first input with a batch dim wins" logic set batch=1 +// and rejected attention_mask as a mismatch; detection now takes 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) + resize_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 broadcast to 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(); + ov::SoPtr inner{std::make_shared(model, m_plugin, m_recorder), {}}; + auto wrapped = ov::npuw::batched::CompiledModel::create(model, m_plugin, inner, /*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"})); +} + } // namespace From 2f1f1ca7d6dacee8e1b7777e33959f15bfd8d7e2 Mon Sep 17 00:00:00 2001 From: Dylan Neve Date: Mon, 29 Jun 2026 11:54:37 +0100 Subject: [PATCH 5/7] Apply clang-format to batched scoring element and tests --- .../src/plugin/npuw/llm_compiled_model.cpp | 6 +++--- .../src/plugin/npuw/v1/elements/batched.cpp | 9 ++++----- .../tests/unit/npuw/batched_element_test.cpp | 18 ++++++------------ .../tests/unit/npuw/llm_test_helpers.hpp | 10 ++++------ 4 files changed, 17 insertions(+), 26 deletions(-) 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 8473a7ee27e2..acc2e70aba7d 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 @@ -10,7 +10,6 @@ #include "llm_compiled_model_utils.hpp" #include "llm_infer_request.hpp" #include "logging.hpp" -#include "v1/elements/batched.hpp" #include "moe_transformations/apply_moe_device_routed_transforms.hpp" #include "npuw_transformations/add_position_ids_param.hpp" #include "npuw_transformations/convert_kvcache_to_precision.hpp" @@ -46,6 +45,7 @@ #include "serialization.hpp" #include "transformations/convert_precision.hpp" #include "util.hpp" +#include "v1/elements/batched.hpp" #include "whisper/prepare_whisper_model.hpp" #include "whisper/whisper_infer_request.hpp" @@ -1628,8 +1628,8 @@ std::shared_ptr ov::npuw::LLMCompiledModel::create_sync_i return non_const_this->create_whisper_infer_request(); } - auto inner = m_is_embedding ? non_const_this->create_embedding_infer_request() - : non_const_this->create_llm_infer_request(); + auto inner = + m_is_embedding ? non_const_this->create_embedding_infer_request() : non_const_this->create_llm_infer_request(); // Batched scoring: wrap the single-sequence request with the batched element so a // [N, ...] input is unrolled into N independent [1, ...] inferences (rows are scored 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 index 1bd8f7f80391..dd3398e41907 100644 --- a/src/plugins/intel_npu/src/plugin/npuw/v1/elements/batched.cpp +++ b/src/plugins/intel_npu/src/plugin/npuw/v1/elements/batched.cpp @@ -28,11 +28,10 @@ bool has_batch_dim(const ov::SoPtr& tensor) { } // namespace -ov::SoPtr ov::npuw::batched::CompiledModel::create( - const std::shared_ptr& model, - const std::shared_ptr& plugin, - ov::SoPtr inner_compiled, - bool enabled) { +ov::SoPtr ov::npuw::batched::CompiledModel::create(const std::shared_ptr& model, + const std::shared_ptr& plugin, + ov::SoPtr inner_compiled, + bool enabled) { OPENVINO_ASSERT(inner_compiled._ptr != nullptr, "Batched compiled model requires an inner compiled model"); // No-op wrapper: hand back the inner model unchanged for the zero-overhead path. 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 index 2b7d38e39794..ed385dae563f 100644 --- a/src/plugins/intel_npu/tests/unit/npuw/batched_element_test.cpp +++ b/src/plugins/intel_npu/tests/unit/npuw/batched_element_test.cpp @@ -60,8 +60,7 @@ class MockState : public ov::IVariableState { std::shared_ptr m_rec; }; -ov::Output port_named(const std::vector>& ports, - const std::string& needle) { +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; @@ -253,8 +252,7 @@ TEST_F(NPUWBatchedElementTest, EachRowScoredIndependentlyAndStacked) { } // 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"})); + EXPECT_EQ(m_recorder->events, (std::vector{"reset", "infer", "reset", "infer", "reset", "infer"})); } TEST_F(NPUWBatchedElementTest, SingleRowScored) { @@ -325,8 +323,7 @@ TEST_F(NPUWBatchedElementTest, SyncInnerConstructorScoresEachRow) { 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"})); + EXPECT_EQ(m_recorder->events, (std::vector{"reset", "infer", "reset", "infer", "reset", "infer"})); } // A shared/broadcast input (batch 1) is fed to every row unsliced -- the slice guard only fires @@ -346,8 +343,7 @@ TEST_F(NPUWBatchedElementTest, BroadcastInputSharedAcrossRows) { 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"})); + 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 @@ -367,8 +363,7 @@ TEST_F(NPUWBatchedElementTest, BatchSizeFromNonLeadingInput) { for (std::size_t r = 0; r < 3; ++r) { EXPECT_FLOAT_EQ(row_value(out, r), 7.0f) << "output row " << r; // input_ids broadcast to every row } - EXPECT_EQ(m_recorder->events, - (std::vector{"reset", "infer", "reset", "infer", "reset", "infer"})); + 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 @@ -393,8 +388,7 @@ TEST_F(NPUWBatchedElementTest, MultipleOutputsStackedIndependently) { 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"})); + EXPECT_EQ(m_recorder->events, (std::vector{"reset", "infer", "reset", "infer", "reset", "infer"})); } } // 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 2f638e024cf9..a92672720351 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}; @@ -129,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); @@ -344,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; }; From d2e03b856acfa11330d55b83ae25677d020ea186 Mon Sep 17 00:00:00 2001 From: Dylan Neve Date: Wed, 8 Jul 2026 13:18:41 +0100 Subject: [PATCH 6/7] NPUW: Compose the batched element at the entry points Rework the batched scoring element per review: - Instantiate batched::CompiledModel at the NPUW entry points -- npuw::ICompiledModel::create() on compilation and the plugin's NPUW import path on blob import -- wrapping the LLMCompiledModel built right there from the same model, instead of LLMCompiledModel wrapping its own infer request from the inside. - Make the element a true npuw::ICompiledModel decorator: inputs() and outputs() forward to the inner model, so the wrapper needs no ov::Model of its own (works on import too) and the request operates on the very same port objects as the inner. - Drop the sync/async inner-request duality: the request holds a single inner IAsyncInferRequest; its infer() runs inline on the calling thread, so no executor nesting is possible in the first place. - Drop init_public_tensors(): instead of allocating host tensors for every port, the request surfaces the inner request's own tensors as the defaults; nothing is allocated up front. - Use ov::npuw::util::view() for row slicing, simplify batch detection (batch is the largest leading dim; leading-dim-1 inputs are shared), and write rows straight into the [N, ...] public output tensors, reusing a caller-bound output that already fits. - Un-serialize the rerank flag: NPUW_TEXT_RERANK is consumed at the entry points, so LLMCompiledModel no longer carries m_is_rerank and the blob format / NPUW_SERIALIZATION_VERSION stay untouched. - Rework the unit tests for the new composition; add coverage for batched::requested(), I/O forwarding and pre-bound output reuse. --- .../src/plugin/npuw/compiled_model.cpp | 10 +- .../src/plugin/npuw/llm_compiled_model.cpp | 32 +- .../src/plugin/npuw/llm_compiled_model.hpp | 1 - .../src/plugin/npuw/v1/elements/batched.cpp | 281 +++++++----------- .../src/plugin/npuw/v1/elements/batched.hpp | 120 ++++---- .../intel_npu/src/plugin/src/plugin.cpp | 9 +- .../tests/unit/npuw/batched_element_test.cpp | 152 +++++----- 7 files changed, 260 insertions(+), 345 deletions(-) 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 66c603c51054..ab27f7347972 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 acc2e70aba7d..a23d073508c5 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 @@ -45,7 +45,6 @@ #include "serialization.hpp" #include "transformations/convert_precision.hpp" #include "util.hpp" -#include "v1/elements/batched.hpp" #include "whisper/prepare_whisper_model.hpp" #include "whisper/whisper_infer_request.hpp" @@ -816,8 +815,10 @@ 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; - auto use_text_rerank_key = pop_option(other_props, std::string("NPUW_TEXT_RERANK")); - m_is_rerank = use_text_rerank_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"); @@ -1623,30 +1624,13 @@ ov::Any ov::npuw::LLMCompiledModel::get_property(const std::string& name) const std::shared_ptr ov::npuw::LLMCompiledModel::create_sync_infer_request() const { auto* non_const_this = const_cast(this); // because of const in API - if (m_is_whisper) { return non_const_this->create_whisper_infer_request(); + } else if (m_is_embedding) { + return non_const_this->create_embedding_infer_request(); + } else { + return non_const_this->create_llm_infer_request(); } - - auto inner = - m_is_embedding ? non_const_this->create_embedding_infer_request() : non_const_this->create_llm_infer_request(); - - // Batched scoring: wrap the single-sequence request with the batched element so a - // [N, ...] input is unrolled into N independent [1, ...] inferences (rows are scored - // one at a time, with the KV-cache reset between them) and stacked back into [N, ...]. - // This is valid only for the non-generating scoring pipelines -- text embedding and - // text reranking -- where rows are independent; the pipeline flags the model as such at - // compile time, so no user-facing option is needed. Generation never sets these. - if (m_is_embedding || m_is_rerank) { - // Drive the single-sequence request synchronously (sync overload) so it runs on the - // calling thread rather than being re-scheduled onto this model's task executor -- the - // outer request is already async on that executor, and nesting on the same executor - // would deadlock. - auto self = std::static_pointer_cast(shared_from_this()); - return std::make_shared(self, std::move(inner)); - } - - return inner; } std::shared_ptr ov::npuw::LLMCompiledModel::create_llm_infer_request() { diff --git a/src/plugins/intel_npu/src/plugin/npuw/llm_compiled_model.hpp b/src/plugins/intel_npu/src/plugin/npuw/llm_compiled_model.hpp index 74fd777eedad..d9796d2d444a 100644 --- a/src/plugins/intel_npu/src/plugin/npuw/llm_compiled_model.hpp +++ b/src/plugins/intel_npu/src/plugin/npuw/llm_compiled_model.hpp @@ -142,7 +142,6 @@ class LLMCompiledModel : public ov::npuw::ICompiledModel { size_t m_decomposed_sdpa_size = 0; bool m_is_embedding = false; - bool m_is_rerank = false; // Create generate model variants with different sizes std::vector> create_generate_model_variants( 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 index dd3398e41907..d849ec93905e 100644 --- a/src/plugins/intel_npu/src/plugin/npuw/v1/elements/batched.cpp +++ b/src/plugins/intel_npu/src/plugin/npuw/v1/elements/batched.cpp @@ -7,51 +7,52 @@ #include #include +#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" -namespace { - -// Zero-copy view of one batch row (axis 0) of a tensor: [N, ...] -> [1, ...]. -ov::SoPtr row_slice(const ov::SoPtr& tensor, std::size_t row) { - ov::Shape start_shape(tensor->get_shape().size(), 0u); - start_shape[0] = row; - ov::Shape end_shape = tensor->get_shape(); - end_shape[0] = row + 1; - return ov::get_tensor_impl(ov::Tensor(ov::make_tensor(tensor), start_shape, end_shape)); -} - -bool has_batch_dim(const ov::SoPtr& tensor) { - return !tensor->get_shape().empty(); +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()); } -} // namespace - -ov::SoPtr ov::npuw::batched::CompiledModel::create(const std::shared_ptr& model, - const std::shared_ptr& plugin, - ov::SoPtr inner_compiled, - bool enabled) { - OPENVINO_ASSERT(inner_compiled._ptr != nullptr, "Batched compiled model requires an inner compiled model"); +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_compiled; + 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"); +} - auto compiled_model = std::make_shared(model, plugin, std::move(inner_compiled)); - return {compiled_model, {}}; +const std::vector>& ov::npuw::batched::CompiledModel::inputs() const { + return m_inner->inputs(); } -ov::npuw::batched::CompiledModel::CompiledModel(const std::shared_ptr& model, - const std::shared_ptr& plugin, - ov::SoPtr inner_compiled) - : ov::ICompiledModel(model, plugin), - m_inner(std::move(inner_compiled)) { - OPENVINO_ASSERT(m_inner._ptr != nullptr, "Batched compiled model requires an inner compiled model"); +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); } @@ -67,6 +68,10 @@ ov::Any ov::npuw::batched::CompiledModel::get_property(const std::string& name) 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(); @@ -74,193 +79,109 @@ std::shared_ptr ov::npuw::batched::CompiledModel::create_ return std::make_shared(self, std::move(inner_request)); } -std::shared_ptr ov::npuw::batched::CompiledModel::create_infer_request() const { - return std::make_shared(create_sync_infer_request(), - get_task_executor(), - get_callback_executor()); -} - ov::npuw::batched::InferRequest::InferRequest(const std::shared_ptr& compiled_model, std::shared_ptr inner_request) : ov::ISyncInferRequest(compiled_model), - m_inner_async(std::move(inner_request)) { - OPENVINO_ASSERT(m_inner_async != nullptr, "Batched element requires a non-null inner request"); - init_public_tensors(); -} + m_inner(std::move(inner_request)) { + OPENVINO_ASSERT(m_inner != nullptr, "Batched element requires a non-null inner request"); -ov::npuw::batched::InferRequest::InferRequest(const std::shared_ptr& compiled_model, - std::shared_ptr inner_request) - : ov::ISyncInferRequest(compiled_model), - m_inner_sync(std::move(inner_request)) { - OPENVINO_ASSERT(m_inner_sync != nullptr, "Batched element requires a non-null inner request"); - init_public_tensors(); -} - -void ov::npuw::batched::InferRequest::init_public_tensors() { - // Allocate a tensor for every public port so get_tensor() never returns an - // uninitialized handle (callers such as the Python infer(dict) dispatcher fetch - // the tensor before populating it). Dynamic dims are sized to 0 and resized later; - // the real [N, ...] tensors are bound by the caller (inputs) or by infer() (outputs). - const auto init_port = [this](const ov::Output& port) { - if (ov::ISyncInferRequest::get_tensor(port)) { - return; - } - const auto& pshape = port.get_partial_shape(); - ov::Shape shape; - if (pshape.is_dynamic()) { - for (const auto& dim : pshape) { - shape.push_back(dim.is_static() ? dim.get_length() : 0); - } - } else { - shape = pshape.to_shape(); - } - set_tensor(port, ov::get_tensor_impl(ov::Tensor(port.get_element_type(), shape))); - }; + // 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()) { - init_port(port); - } - for (const auto& port : get_outputs()) { - init_port(port); - } -} - -const std::vector>& ov::npuw::batched::InferRequest::inner_inputs() const { - return m_inner_sync ? m_inner_sync->get_inputs() : m_inner_async->get_inputs(); -} - -const std::vector>& ov::npuw::batched::InferRequest::inner_outputs() const { - return m_inner_sync ? m_inner_sync->get_outputs() : m_inner_async->get_outputs(); -} - -void ov::npuw::batched::InferRequest::inner_set_tensor(const ov::Output& port, - const ov::SoPtr& tensor) { - if (m_inner_sync) { - m_inner_sync->set_tensor(port, tensor); - } else { - m_inner_async->set_tensor(port, tensor); - } -} - -ov::SoPtr ov::npuw::batched::InferRequest::inner_get_tensor(const ov::Output& port) const { - return m_inner_sync ? m_inner_sync->get_tensor(port) : m_inner_async->get_tensor(port); -} - -void ov::npuw::batched::InferRequest::inner_infer() { - if (m_inner_sync) { - m_inner_sync->infer(); - } else { - m_inner_async->infer(); + if (auto tensor = m_inner->get_tensor(port)) { + set_tensor(port, tensor); + } } } -std::vector> ov::npuw::batched::InferRequest::inner_query_state() const { - return m_inner_sync ? m_inner_sync->query_state() : m_inner_async->query_state(); -} - void ov::npuw::batched::InferRequest::infer() { std::lock_guard lock(m_mutex); - const auto& wrapper_inputs = get_inputs(); - const auto& wrapper_outputs = get_outputs(); - const auto& in_ports = inner_inputs(); - const auto& out_ports = inner_outputs(); - - OPENVINO_ASSERT(wrapper_inputs.size() == in_ports.size() && wrapper_outputs.size() == out_ports.size(), - "Batched element: inner request I/O does not match the wrapped model"); + const auto& in_ports = get_inputs(); + const auto& out_ports = get_outputs(); + OPENVINO_ASSERT(!in_ports.empty(), "Batched element: the wrapped model has no inputs"); - // Snapshot the public inputs once -- read repeatedly below. + // Snapshot the public inputs. The batch is the largest leading dimension across + // them: every input must either carry the batch ([N, ...], sliced per row) or be + // shared across rows ([1, ...], bound whole). std::vector> in_tensors; - in_tensors.reserve(wrapper_inputs.size()); - for (const auto& port : wrapper_inputs) { - in_tensors.push_back(get_tensor(port)); - } - - // Batch is the largest leading dim across inputs, so a shared [1, ...] input is broadcast - // rather than mistaken for the batch. Stays 1 when nothing is batched. + in_tensors.reserve(in_ports.size()); std::size_t batch = 1; - bool any_batched = false; - for (const auto& tensor : in_tensors) { - if (!has_batch_dim(tensor)) { - continue; - } - const std::size_t in_batch = tensor->get_shape()[0]; - batch = any_batched ? std::max(batch, in_batch) : in_batch; - any_batched = true; + 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"); + batch = std::max(batch, shape[0]); + in_tensors.push_back(std::move(tensor)); } - // A zero-sized batch has no rows and would publish unpopulated outputs; reject it. - OPENVINO_ASSERT(batch > 0, "Batched element: batch size must be > 0, got an input with batch dimension 0."); - - // Every batched input must match `batch` or be a broadcast (dim 0 == 1); anything else - // cannot be sliced per row. - for (std::size_t i = 0; i < wrapper_inputs.size(); ++i) { - if (!has_batch_dim(in_tensors[i])) { - continue; - } + for (std::size_t i = 0; i < in_ports.size(); ++i) { const std::size_t in_batch = in_tensors[i]->get_shape()[0]; OPENVINO_ASSERT(in_batch == batch || in_batch == 1, "Batched element: input '", - wrapper_inputs[i].get_any_name(), + in_ports[i].get_any_name(), "' has batch dimension ", in_batch, " which is neither the inferred batch size ", batch, - " nor 1 (broadcast)."); + " nor 1 (shared)."); } - // States are stable across rows: query once, reset per row so row i never sees row i-1. - const auto inner_states = inner_query_state(); - const auto reset_inner_state = [&] { + // 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) { for (const auto& state : inner_states) { state->reset(); } - }; - - // Slice per-row inputs out of [batch, ...]; pass broadcast/non-batched inputs through whole. - const auto bind_row = [&](std::size_t row) { - for (std::size_t i = 0; i < wrapper_inputs.size(); ++i) { + for (std::size_t i = 0; i < in_ports.size(); ++i) { const auto& full = in_tensors[i]; - const bool sliceable = batch > 1 && has_batch_dim(full) && full->get_shape()[0] == batch; - inner_set_tensor(in_ports[i], sliceable ? row_slice(full, row) : full); + m_inner->set_tensor(in_ports[i], full->get_shape()[0] == 1 ? full : ov::npuw::util::view(full, 0, row, 1)); } - }; - - // Per-row outputs stacked along axis 0. A single row has nothing to stack, so its output is - // published as-is (no aggregation buffer or copy); a non-batched output collapses to the last row. - std::vector> aggregated_outputs(wrapper_outputs.size()); + m_inner->infer(); - for (std::size_t row = 0; row < batch; ++row) { - reset_inner_state(); - bind_row(row); - inner_infer(); - - for (std::size_t i = 0; i < wrapper_outputs.size(); ++i) { - const auto inner_out = inner_get_tensor(out_ports[i]); - if (batch == 1) { - aggregated_outputs[i] = inner_out; - continue; - } - if (!aggregated_outputs[i]) { - ov::Shape out_shape = inner_out->get_shape(); - if (!out_shape.empty()) { - out_shape[0] = batch; - } - aggregated_outputs[i] = ov::get_tensor_impl(ov::Tensor(inner_out->get_element_type(), out_shape)); - } - if (has_batch_dim(inner_out)) { - inner_out->copy_to(row_slice(aggregated_outputs[i], row)._ptr); - } else { - inner_out->copy_to(aggregated_outputs[i]._ptr); - } + 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); + } + for (const auto& port : out_ports) { + m_inner->get_tensor(port)->copy_to(ov::npuw::util::view(get_tensor(port), 0, row, 1)._ptr); } } +} - for (std::size_t i = 0; i < wrapper_outputs.size(); ++i) { - set_tensor(wrapper_outputs[i], aggregated_outputs[i]); +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: batched tensors are unrolled and validated per row by the inner request. + // 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 { @@ -270,5 +191,5 @@ std::vector> ov::npuw::batched::InferRequest::quer } std::vector ov::npuw::batched::InferRequest::get_profiling_info() const { - return {}; + 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 index 204f551ce0f8..ee0206ed190f 100644 --- a/src/plugins/intel_npu/src/plugin/npuw/v1/elements/batched.hpp +++ b/src/plugins/intel_npu/src/plugin/npuw/v1/elements/batched.hpp @@ -8,50 +8,56 @@ #include #include +#include "../../compiled_model.hpp" #include "openvino/runtime/iasync_infer_request.hpp" -#include "openvino/runtime/icompiled_model.hpp" #include "openvino/runtime/isync_infer_request.hpp" #include "openvino/runtime/so_ptr.hpp" namespace ov::npuw::batched { -class InferRequest; +// 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 wrapper that adds batched (batch > 1) execution on top of an +// 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 ov::ICompiledModel decorator that exposes the same I/O as the -// model it wraps and forwards property queries 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 stacked -// back into [N, ...] tensors along the batch dimension (axis 0). +// 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 rows are not independent. +// identical and batching is purely a throughput/ergonomics choice. It is NOT +// valid for autoregressive generation, where state must persist across calls. // -// It composes on top of the other elements, e.g.: -// -// Batched( AccuracyChecked( Failsafe(NPU -> CPU) ) ) -// -// and the inner request needs no awareness of batching at all -- a stock -// single-sequence LLM/embedding infer request works unchanged. -class CompiledModel final : public ov::ICompiledModel { +// 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_compiled unwrapped when enabled == false, - // keeping the zero-overhead path trivial (mirrors accuracy_checked::create). - static ov::SoPtr create(const std::shared_ptr& model, - const std::shared_ptr& plugin, - ov::SoPtr inner_compiled, - bool enabled); + // 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); - CompiledModel(const std::shared_ptr& model, - const std::shared_ptr& plugin, - ov::SoPtr inner_compiled); + // 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; @@ -59,46 +65,30 @@ class CompiledModel final : public ov::ICompiledModel { 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; - std::shared_ptr create_infer_request() const override; private: - friend class InferRequest; - - ov::SoPtr m_inner; + std::shared_ptr m_inner; }; -// Sync infer request that unrolls a batched inference over a single-sequence +// Sync infer request that unrolls a batched inference over the single-sequence // inner request. // -// On each infer() call it takes the batch size N as the largest leading dimension -// across the inputs (so a shared [1, ...] input is broadcast, not mistaken for the -// batch), then for each row 0..N-1: resets the inner request's variable state, binds -// the row's [1, ...] slice of every per-row input, runs the inner request, and copies -// the inner [1, ...] output into row i of the wrapper's [N, ...] output tensors. -// N == 1 publishes the inner outputs directly. The public input/output tensors are -// held by the ISyncInferRequest base; only the inner request sees [1, ...] shapes. -// -// It is constructed from the compiled model whose I/O it exposes plus the inner -// request to drive. This lets it be produced both by Batched::CompiledModel -// (standalone, composable element) and reused directly by a pipeline that -// already owns a single-sequence request (e.g. NPUW's LLMCompiledModel wrapping -// its LLMInferRequest for scoring), without that pipeline having to wrap its -// whole compiled model. +// 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: - // Drive an async inner request (used by the standalone Batched::CompiledModel decorator, - // where the inner is a separate compiled model with its own task executor). InferRequest(const std::shared_ptr& compiled_model, std::shared_ptr inner_request); - // Drive a sync inner request directly, on the calling thread, without an executor. Used when a - // pipeline wraps its own single-sequence request (e.g. NPUW's LLMCompiledModel wrapping its - // LLMInferRequest): this avoids a nested-executor deadlock that would otherwise occur if the - // inner ran on the same task executor as the outer (async) request. - InferRequest(const std::shared_ptr& compiled_model, - std::shared_ptr inner_request); - void infer() override; void check_tensors() const override; @@ -106,19 +96,13 @@ class InferRequest final : public ov::ISyncInferRequest { std::vector get_profiling_info() const override; private: - void init_public_tensors(); - - // The inner single-sequence request, driven row by row. Exactly one of the two handles is set; - // the small accessors below hide which interface (sync/async) is in use from infer(). - const std::vector>& inner_inputs() const; - const std::vector>& inner_outputs() const; - void inner_set_tensor(const ov::Output& port, const ov::SoPtr& tensor); - ov::SoPtr inner_get_tensor(const ov::Output& port) const; - void inner_infer(); - std::vector> inner_query_state() const; - - std::shared_ptr m_inner_async; - std::shared_ptr m_inner_sync; + // 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; }; diff --git a/src/plugins/intel_npu/src/plugin/src/plugin.cpp b/src/plugins/intel_npu/src/plugin/src/plugin.cpp index fcc4a43dd2fb..9bbe8090e118 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/npuw/batched_element_test.cpp b/src/plugins/intel_npu/tests/unit/npuw/batched_element_test.cpp index ed385dae563f..8a51ae4c91a0 100644 --- a/src/plugins/intel_npu/tests/unit/npuw/batched_element_test.cpp +++ b/src/plugins/intel_npu/tests/unit/npuw/batched_element_test.cpp @@ -4,16 +4,18 @@ // 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 stacks the per-row outputs back to [N, ...]. -// 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. +// 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 @@ -133,22 +135,17 @@ class MockInnerSync : public ov::ISyncInferRequest { std::shared_ptr m_state; }; -class MockInnerCompiled : public ov::ICompiledModel { +class MockInnerCompiled : public ov::npuw::ICompiledModel { public: MockInnerCompiled(const std::shared_ptr& model, const std::shared_ptr& plugin, std::shared_ptr rec) - : ov::ICompiledModel(model, plugin), + : 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); } - std::shared_ptr create_infer_request() const override { - return std::make_shared(create_sync_infer_request(), - get_task_executor(), - get_callback_executor()); - } void export_model(std::ostream&) const override {} std::shared_ptr get_runtime_model() const override { return nullptr; @@ -170,22 +167,25 @@ class NPUWBatchedElementTest : public ::testing::Test { m_recorder = std::make_shared(); } - ov::SoPtr make_inner() { - return {std::make_shared(m_model, m_plugin, m_recorder), {}}; + std::shared_ptr make_inner() { + return std::make_shared(m_model, m_plugin, m_recorder); } - ov::SoPtr wrap(bool enabled) { - return ov::npuw::batched::CompiledModel::create(m_model, m_plugin, make_inner(), enabled); + std::shared_ptr wrap(bool enabled) { + return ov::npuw::batched::CompiledModel::create(make_inner(), m_plugin, enabled); } - // Resize a named input in place and return it. Templated so it serves both the async wrapper - // and a bare sync request -- both expose get_inputs()/get_tensor(). + // Bind a fresh zero-filled tensor of the port's element type to a named input. template - static ov::SoPtr resize_input(const std::shared_ptr& req, - const std::string& name, - const ov::Shape& shape) { - const auto tensor = req->get_tensor(port_named(req->get_inputs(), name)); - tensor->set_shape(shape); + 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; } @@ -194,7 +194,7 @@ class NPUWBatchedElementTest : public ::testing::Test { 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(resize_input(req, "input_ids", {batch, len})->data()); + 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); } @@ -206,9 +206,9 @@ class NPUWBatchedElementTest : public ::testing::Test { const std::size_t batch = ids.size(); const std::size_t len = ids.empty() ? 0 : ids.front().size(); set_input_ids(req, ids); - resize_input(req, "attention_mask", {batch, len}); - resize_input(req, "position_ids", {batch, len}); - resize_input(req, "beam_idx", {batch}); + 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. @@ -222,17 +222,28 @@ class NPUWBatchedElementTest : public ::testing::Test { 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(m_model, m_plugin, inner, /*enabled=*/false); - EXPECT_EQ(wrapped._ptr, inner._ptr); + 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(m_model, m_plugin, inner, /*enabled=*/true); - EXPECT_NE(wrapped._ptr, inner._ptr); - EXPECT_NE(std::dynamic_pointer_cast(wrapped._ptr), nullptr); + 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) { @@ -270,18 +281,19 @@ TEST_F(NPUWBatchedElementTest, SingleRowScored) { EXPECT_EQ(m_recorder->events, (std::vector{"reset", "infer"})); } -// A zero-row batch has nothing to score and would publish an unpopulated output. +// 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(); - resize_input(req, "input_ids", {0, 4}); + 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("batch size must be > 0"), std::string::npos); + EXPECT_NE(error_message(ex).find("zero-sized batch dimension"), std::string::npos); } EXPECT_TRUE(m_recorder->events.empty()); } @@ -293,7 +305,7 @@ TEST_F(NPUWBatchedElementTest, MismatchedBatchRejected) { auto req = wrapped->create_infer_request(); bind_inputs(req, {{1, 2}, {3, 4}, {5, 6}}); - resize_input(req, "attention_mask", {2, 2}); + bind_input(req, "attention_mask", {2, 2}); try { req->infer(); @@ -304,37 +316,15 @@ TEST_F(NPUWBatchedElementTest, MismatchedBatchRejected) { EXPECT_TRUE(m_recorder->events.empty()); } -// The production LLM/rerank wiring builds the element from a *sync* inner request (driven on the -// calling thread to dodge a nested-executor deadlock), bypassing batched::CompiledModel. Exercise -// that constructor directly. -TEST_F(NPUWBatchedElementTest, SyncInnerConstructorScoresEachRow) { - auto inner_compiled = std::make_shared(m_model, m_plugin, m_recorder); - std::shared_ptr inner_sync = inner_compiled->create_sync_infer_request(); - std::shared_ptr compiled = inner_compiled; - auto req = std::make_shared(compiled, std::move(inner_sync)); - - 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(req->get_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"})); -} - -// A shared/broadcast input (batch 1) is fed to every row unsliced -- the slice guard only fires -// when the leading dim equals the batch, so a [1, L] input is never sliced out of range. -TEST_F(NPUWBatchedElementTest, BroadcastInputSharedAcrossRows) { +// 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); - resize_input(req, "attention_mask", {1, 4}); // shared across the 3 rows + bind_input(req, "attention_mask", {1, 4}); // shared across the 3 rows req->infer(); @@ -346,22 +336,21 @@ TEST_F(NPUWBatchedElementTest, BroadcastInputSharedAcrossRows) { 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 old "first input with a batch dim wins" logic set batch=1 -// and rejected attention_mask as a mismatch; detection now takes the largest leading dim. +// 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) - resize_input(req, "attention_mask", {3, 4}); // real batch N = 3 on a non-leading input + 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 broadcast to every row + 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"})); } @@ -370,8 +359,8 @@ TEST_F(NPUWBatchedElementTest, BatchSizeFromNonLeadingInput) { // two outputs must differ per row. TEST_F(NPUWBatchedElementTest, MultipleOutputsStackedIndependently) { auto model = build_two_output_model(); - ov::SoPtr inner{std::make_shared(model, m_plugin, m_recorder), {}}; - auto wrapped = ov::npuw::batched::CompiledModel::create(model, m_plugin, inner, /*enabled=*/true); + 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}}; @@ -391,4 +380,27 @@ TEST_F(NPUWBatchedElementTest, MultipleOutputsStackedIndependently) { 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 From 7797220962b4efe126a8cbc942aed95305a0a5d4 Mon Sep 17 00:00:00 2001 From: Dylan Neve Date: Thu, 16 Jul 2026 11:02:43 +0100 Subject: [PATCH 7/7] NPUW: Extract the batch snapshot into its own function, meter the unroll Per review, pull the input snapshot / batch-size derivation out of InferRequest::infer() into a dedicated extract_batch(), returning the snapshotted tensors and the derived batch as a BatchedInputs struct. infer() is left as just the unroll: bind row, infer, copy row out. No behavioural change -- the validation and the "largest leading dim, leading-dim-1 is shared" rule are carried over verbatim. Also meter the unroll phases with the existing ov::npuw::perf::Profile (reported on request destruction under OPENVINO_NPUW_PROF=1 on a developer build), so extract_batch()'s cost can be read against the inference it wraps rather than guessed at. Measured on the deviceless unit-test harness (synthetic Qwen-style reranker, 4 inputs, seq len 512), calling extract_batch() directly over 200k iterations: batch=1 -> 239 ns/call batch=8 -> 239 ns/call batch=16 -> 240 ns/call batch=64 -> 241 ns/call It runs once per infer() and is flat in batch size and in tensor data size, as expected: the work is O(number of inputs) -- a get_tensor() lookup and a shape read per input, ~60 ns each -- and never touches tensor data. Against a real batched scoring call, where each of the N rows costs a full inner model inference in the millisecond range, the snapshot is a ~1e-5 fraction of the total and does not warrant caching across calls (which would have to be invalidated on every set_tensor() anyway). --- .../src/plugin/npuw/v1/elements/batched.cpp | 70 ++++++++++++------- .../src/plugin/npuw/v1/elements/batched.hpp | 18 +++++ 2 files changed, 62 insertions(+), 26 deletions(-) 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 index d849ec93905e..8d52122eefdf 100644 --- a/src/plugins/intel_npu/src/plugin/npuw/v1/elements/batched.cpp +++ b/src/plugins/intel_npu/src/plugin/npuw/v1/elements/batched.cpp @@ -7,6 +7,7 @@ #include #include +#include "../../logging.hpp" #include "../../util.hpp" #include "intel_npu/npuw_private_properties.hpp" #include "openvino/core/except.hpp" @@ -85,6 +86,9 @@ ov::npuw::batched::InferRequest::InferRequest(const std::shared_ptr lock(m_mutex); - +ov::npuw::batched::InferRequest::BatchedInputs ov::npuw::batched::InferRequest::extract_batch() const { const auto& in_ports = get_inputs(); - const auto& out_ports = get_outputs(); OPENVINO_ASSERT(!in_ports.empty(), "Batched element: the wrapped model has no inputs"); - // Snapshot the public inputs. The batch is the largest leading dimension across - // them: every input must either carry the batch ([N, ...], sliced per row) or be - // shared across rows ([1, ...], bound whole). - std::vector> in_tensors; - in_tensors.reserve(in_ports.size()); - std::size_t batch = 1; + 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(), "'"); @@ -121,20 +118,34 @@ void ov::npuw::batched::InferRequest::infer() { "Batched element: input '", port.get_any_name(), "' has a zero-sized batch dimension - batch size must be > 0"); - batch = std::max(batch, shape[0]); - in_tensors.push_back(std::move(tensor)); + 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 = in_tensors[i]->get_shape()[0]; - OPENVINO_ASSERT(in_batch == batch || in_batch == 1, + 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 ", - batch, + 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 @@ -142,23 +153,30 @@ void ov::npuw::batched::InferRequest::infer() { // [N, ...] public output tensors. const auto inner_states = m_inner->query_state(); for (std::size_t row = 0; row < batch; ++row) { - for (const auto& state : inner_states) { - state->reset(); - } - for (std::size_t i = 0; i < in_ports.size(); ++i) { - const auto& full = in_tensors[i]; - m_inner->set_tensor(in_ports[i], full->get_shape()[0] == 1 ? full : ov::npuw::util::view(full, 0, row, 1)); - } - m_inner->infer(); + 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); } - for (const auto& port : out_ports) { - m_inner->get_tensor(port)->copy_to(ov::npuw::util::view(get_tensor(port), 0, row, 1)._ptr); - } + 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); + } + }); } } 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 index ee0206ed190f..f1cbec31588a 100644 --- a/src/plugins/intel_npu/src/plugin/npuw/v1/elements/batched.hpp +++ b/src/plugins/intel_npu/src/plugin/npuw/v1/elements/batched.hpp @@ -9,6 +9,7 @@ #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" @@ -96,6 +97,19 @@ class InferRequest final : public ov::ISyncInferRequest { 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 @@ -104,6 +118,10 @@ class InferRequest final : public ov::ISyncInferRequest { 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