NPUW: Batched scoring as a composable element (reranker / embedding) by dylanneve1 · Pull Request #36375 · openvinotoolkit/openvino

dylanneve1 · 2026-06-12T09:14:18Z

Details

NPUW compiles every LLM-family graph with a static batch size of 1 (reshape_to_static pins the batch dim to 1 on all inputs and the KV-cache). As a result, a batched request that scores N independent inputs in a single inference — e.g. GenAI's TextRerankPipeline scoring N documents, or a text-embedding model — fails for N>1, while N==1 works.

This PR adds batched scoring support without changing the compiled graph, as a composable NPUW element: a runtime decorator (the same pattern as failsafe and accuracy_checked under v1/elements/) that wraps the unchanged single-sequence request and unrolls the batch row by row.

On each batched infer() the element:

resets the inner request's variable state (KV-cache) so each row is an independent prompt;
binds row n's [1, L] slice of every input (zero-copy ROI view);
runs the unchanged batch-1 prefill;
copies the row's output into row n of an aggregated [N, ...] output tensor.

Because attention never crosses batch rows, batched and per-row scoring are the same arithmetic — the result is bit-for-bit identical to scoring each row separately. The element therefore targets single-shot scoring/embedding pipelines whose rows are independent; it resets state between rows and is not valid for autoregressive generation, so it is strictly opt-in.

What's in this PR

v1/elements/batched.{hpp,cpp} — the batched::CompiledModel / batched::InferRequest decorator. create() returns the inner unwrapped when disabled (zero-overhead path), and composes on top of the other elements, e.g. Batched(AccuracyChecked(Failsafe(NPU->CPU))). The inner request needs no awareness of batching.
NPUW_LLM_BATCH_UNROLL option (default off) and wiring in LLMCompiledModel::create_sync_infer_request: when set, the single-sequence LLM/embedding request is wrapped with the element. Driving the inner request synchronously avoids a nested-executor deadlock (the outer request is already async on this model's task executor).
Unit tests (batched_element_test.cpp): a recording mock inner verifies that a batched infer matches per-row scoring, that state is reset between rows, that each row's input slice is fed correctly, and that outputs are stacked into the correct slots; plus the zero-overhead-when-disabled and single-row-transparent paths.

Why an element (and not a compiled batch dimension)

An element is a runtime decorator over an unchanged inner request, so it keeps the batch-1 graph and the whole NPUW pipeline untouched — minimal blast radius, and one element transparently covers the reranker, the embedding pipeline, and any future single-shot scorer. (Compiling a true static batch-B graph is a separate, heavier direction with its own trade-offs; this PR is the lightweight, reusable path.)

Validation

New NPUWBatchedElementTest unit tests pass; full ov_npu_unit_tests suite green (the batch-1 default path is unaffected).
Validated end-to-end on the real Qwen3-Reranker-0.6B model through the full NPUW pipeline with NPUW_DEVICES=CPU: a batched [N, L] scoring infer is bit-for-bit identical to per-row scoring (max_abs_diff = 0.0).
Confirmed on real NPU hardware (LNL) end-to-end via GenAI's TextRerankPipeline: multiple documents reranked in one call, deterministic and correctly ordered.

Tickets

N/A

Checklist

Tests added/updated
No change to the default (batch-1) behaviour
Validated on real NPU hardware (reranker, LNL)

Copilot

Pull request overview

Adds an NPUW v1 “batched” element that enables batch>1 scoring/embedding by unrolling a batched infer into per-row batch-1 inner inferences (resetting variable state between rows) and stacking outputs back into batched tensors. This is wired into LLMCompiledModel::create_sync_infer_request() behind a new opt-in NPUW_LLM_BATCH config option, and covered by new unit tests.

Changes:

Introduce ov::npuw::batched::{CompiledModel, InferRequest} decorator implementing row-wise batch unrolling for batch-1 inner requests.
Add NPUW_LLM_BATCH option and wrap LLM/embedding sync infer requests with the batched element when enabled.
Add unit tests for the batched element and wire the new source into the unit test build.

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
src/plugins/intel_npu/tests/unit/npuw/batched_element_test.cpp	Adds unit coverage for batched element behavior (per-row slicing, state reset, output stacking, disabled path).
src/plugins/intel_npu/tests/unit/CMakeLists.txt	Links `batched.cpp` into the unit test target.
src/plugins/intel_npu/src/plugin/npuw/v1/elements/batched.hpp	Declares the new batched element compiled model + infer request decorator API.
src/plugins/intel_npu/src/plugin/npuw/v1/elements/batched.cpp	Implements batched unrolling, ROI slicing, per-row state reset, and output aggregation.
src/plugins/intel_npu/src/plugin/npuw/llm_compiled_model.cpp	Adds the `NPUW_LLM_BATCH` opt-in wrapping for sync infer request creation and exposes the property binding.
src/plugins/intel_npu/src/al/include/intel_npu/config/npuw_option_defs.inc	Defines the new `NPUW_LLM_BATCH` configuration option (default off).

Copilot

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.

On NPU a reranker runs through NPUW, which compiles the model with a static batch dimension of 1. A single rerank() call scoring N documents therefore could not be served by one batched infer. Add utils::update_npu_config_text_rerank (mirroring the text embedding NPU config helper) and call it from TextRerankPipeline when the model targets NPU. It sets NPU_USE_NPUW/NPUW_LLM plus NPUW_TEXT_RERANK, which tags the model as a non-generating scoring pipeline so the plugin wraps the request in its batched-scoring element: a [N, L] batch is unrolled into N independent [1, L] inferences (resetting the KV-cache between rows) and stacked back into [N, ...]. Rows are scored independently, so the aggregated output matches per-row scoring. Unlike embedding, a reranker is a standard causal LM (no prefill-only rebuild), and the batched element derives the batch size from the input tensor, so no batch/seq-len/prompt-length overrides are set -- only the rerank tag. Requires the NPUW NPUW_TEXT_RERANK option (openvinotoolkit/openvino#36375).

dmatveev · 2026-07-07T22:42:37Z

+bool has_batch_dim(const ov::SoPtr<ov::ITensor>& tensor) {
+    return !tensor->get_shape().empty();
+}


I am not sure if that's the right way to detect this

dmatveev · 2026-07-07T22:44:57Z

+    init_public_tensors();
+}
+
+void ov::npuw::batched::InferRequest::init_public_tensors() {


not sure if it is right

dmatveev · 2026-07-07T22:45:22Z

+        } else {
+            shape = pshape.to_shape();
+        }
+        set_tensor(port, ov::get_tensor_impl(ov::Tensor(port.get_element_type(), shape)));


these are CPU tensors by default. Also, why should we do this?

dmatveev · 2026-07-07T22:47:13Z

+    // 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;
+    }


This sounds overcomplicated

dmatveev · 2026-07-07T22:49:05Z

+    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);
+            }
+        }
+    }
+


not sure what this handling is for here, too.

why do we need aggregated outputs? Don't you produce [N,...] tensor for [N,...] input?

dmatveev · 2026-07-07T22:50:05Z

            m_kvcache_desc.dim & m_kvcache_desc.max_generation_token_len & m_kvcache_desc.v_tensors_transposed_pre &
            m_kvcache_desc.v_tensors_transposed_gen & m_prefill_chunk_size & m_use_chunk_prefill & m_max_lora_rank &
            m_enable_prefix_caching & m_prefix_caching_block_size & m_prefix_caching_max_num_blocks & m_is_whisper &
-            m_eos_token_id & m_decomposed_sdpa_size & m_is_eagle & m_is_embedding & m_is_block_kv_cache;
+            m_eos_token_id & m_decomposed_sdpa_size & m_is_eagle & m_is_embedding & m_is_block_kv_cache & m_is_rerank;


we'll never make this compatible :(

dmatveev · 2026-07-07T22:52:36Z

+        auto self = std::static_pointer_cast<const ov::ICompiledModel>(shared_from_this());
+        return std::make_shared<ov::npuw::batched::InferRequest>(self, std::move(inner));


this is strange. Batched is supposed to wrap the LLMCompiledModel for this particular rerank case. Why we have LLMCompiledModel wrapping itself here instead? This order doesn't look correct.

dmatveev · 2026-07-07T22:54:33Z

    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<bool>() == 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<bool>() == true;


I don't see where the Batched::CompiledModel is instantiated. I suppose it should be at npuw::CompiledModel::create as that's the entry point.

shampson-intel · 2026-07-15T10:58:11Z

Tests passed, currently running smoke test

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.

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).

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.

…le-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.

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.

AlexanderKalistratov

Is batch size decided at model compilation time or it is purely dynamic and depends on inputs batch?

AlexanderKalistratov · 2026-07-16T08:04:15Z

+std::shared_ptr<ov::npuw::ICompiledModel> ov::npuw::batched::CompiledModel::create(
+    const std::shared_ptr<ov::npuw::ICompiledModel>& inner,
+    const std::shared_ptr<const ov::IPlugin>& plugin,
+    bool enabled) {


enabled as a parameter feels weird. Since it is passed it can be handled from outside more explicitly.

AlexanderKalistratov · 2026-07-16T09:18:56Z

+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);


So... I don't think it would actually work. Even though it is just a runtime wrapper you must recreate it from somewhere to create it when you would deserialize. And I don't see where it is happens.

Need to figure out how to handle that properly.

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).

github-actions Bot added category: NPU OpenVINO NPU plugin category: NPUW NPUW plugin labels Jun 12, 2026

dylanneve1 force-pushed the dneve/npuw-llm-batch-unroll branch from f52ab66 to 5f0c2f0 Compare June 16, 2026 09:31

dylanneve1 changed the title ~~NPUW: Support batched prefill scoring via row-by-row unrolling~~ NPUW: Batched scoring as a composable element (reranker / embedding) Jun 16, 2026

github-actions Bot added the category: build OpenVINO cmake script / infra label Jun 16, 2026

dylanneve1 requested a review from Copilot June 16, 2026 13:08

Copilot started reviewing on behalf of dylanneve1 June 16, 2026 13:09 View session

Copilot AI reviewed Jun 16, 2026

View reviewed changes

Comment thread src/plugins/intel_npu/src/plugin/npuw/v1/elements/batched.cpp Outdated

Comment thread src/plugins/intel_npu/src/plugin/npuw/v1/elements/batched.cpp Outdated

Comment thread src/plugins/intel_npu/src/plugin/npuw/v1/elements/batched.cpp Outdated

dylanneve1 force-pushed the dneve/npuw-llm-batch-unroll branch 2 times, most recently from dfebddb to b126060 Compare June 16, 2026 15:18

github-actions Bot removed the category: build OpenVINO cmake script / infra label Jun 16, 2026

dylanneve1 requested a review from Copilot June 17, 2026 08:45

Copilot started reviewing on behalf of dylanneve1 June 17, 2026 08:46 View session

Copilot AI reviewed Jun 17, 2026

View reviewed changes

Comment thread src/plugins/intel_npu/src/plugin/npuw/v1/elements/batched.cpp Outdated

Comment thread src/plugins/intel_npu/src/al/include/intel_npu/config/npuw_option_defs.inc Outdated

dylanneve1 force-pushed the dneve/npuw-llm-batch-unroll branch from b126060 to f89467d Compare June 17, 2026 10:56

dylanneve1 force-pushed the dneve/npuw-llm-batch-unroll branch from f89467d to 408eb96 Compare June 25, 2026 15:28

github-actions Bot added the category: build OpenVINO cmake script / infra label Jun 25, 2026

dylanneve1 force-pushed the dneve/npuw-llm-batch-unroll branch 2 times, most recently from a84e44c to fa76732 Compare June 26, 2026 10:13

dylanneve1 marked this pull request as ready for review June 29, 2026 10:15

dylanneve1 requested review from a team as code owners June 29, 2026 10:15

dylanneve1 force-pushed the dneve/npuw-llm-batch-unroll branch from 275d1e5 to 007c23c Compare June 30, 2026 08:26

dylanneve1 mentioned this pull request Jul 1, 2026

[RAG] Enable batched text reranking on NPU via NPUW_TEXT_RERANK openvinotoolkit/openvino.genai#4064

Open

dylanneve1 force-pushed the dneve/npuw-llm-batch-unroll branch from 087550c to 1de0a50 Compare July 2, 2026 10:40

dylanneve1 force-pushed the dneve/npuw-llm-batch-unroll branch from 1de0a50 to efaf895 Compare July 3, 2026 09:45

dylanneve1 requested review from AsyaPronina and dmatveev July 5, 2026 11:03

dylanneve1 requested a review from AlexanderKalistratov July 5, 2026 11:03

dylanneve1 force-pushed the dneve/npuw-llm-batch-unroll branch 2 times, most recently from 42450a3 to 3d557bb Compare July 6, 2026 15:50

dylanneve1 force-pushed the dneve/npuw-llm-batch-unroll branch 2 times, most recently from 08d8b3c to b43cad7 Compare July 7, 2026 14:21

dmatveev reviewed Jul 7, 2026

View reviewed changes

dylanneve1 force-pushed the dneve/npuw-llm-batch-unroll branch 2 times, most recently from 3ec8588 to 0c266a9 Compare July 8, 2026 14:00

dylanneve1 added 6 commits July 15, 2026 13:41

Apply clang-format to batched scoring element and tests

2f1f1ca

dylanneve1 force-pushed the dneve/npuw-llm-batch-unroll branch from 0c266a9 to d2e03b8 Compare July 15, 2026 12:41

AlexanderKalistratov reviewed Jul 16, 2026

View reviewed changes

		auto self = std::static_pointer_cast<const ov::ICompiledModel>(shared_from_this());
		return std::make_shared<ov::npuw::batched::InferRequest>(self, std::move(inner));

Uh oh!

Conversation

dylanneve1 commented Jun 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Details

What's in this PR

Why an element (and not a compiled batch dimension)

Validation

Tickets

Checklist

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

shampson-intel commented Jul 15, 2026

Uh oh!

AlexanderKalistratov left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

dylanneve1 commented Jun 12, 2026 •

edited

Loading