diff --git a/src/plugins/intel_npu/src/plugin/npuw/whisper/prepare_whisper_model.cpp b/src/plugins/intel_npu/src/plugin/npuw/whisper/prepare_whisper_model.cpp index 49409b1c863937..35706ab086cbcb 100644 --- a/src/plugins/intel_npu/src/plugin/npuw/whisper/prepare_whisper_model.cpp +++ b/src/plugins/intel_npu/src/plugin/npuw/whisper/prepare_whisper_model.cpp @@ -4,6 +4,7 @@ #include "prepare_whisper_model.hpp" +#include #include #include "../llm_compiled_model_utils.hpp" @@ -101,6 +102,33 @@ class AttentionMaskInputPast_2 : public ov::pass::MatcherPass { } }; +// If cross-attention SDPA has already been decomposed (by GenAI, or by +// decompose_scaled_dot_product_attention_for_whisper() below), there's no SDPA node left +// to find for it. The decomposition always tags the "QK scaled scores" node with this +// well-known tensor name, so it doubles as a reliable marker for where the block lives. +std::vector> find_decomposed_cross_attn_score_nodes( + const std::shared_ptr& model) { + std::vector> found; + for (const auto& op : model->get_ordered_ops()) { + bool matched = false; + for (const auto& output : op->outputs()) { + for (const auto& name : output.get_names()) { + if (name.find("cross_attention_qk_scaled_scores") != std::string::npos) { + matched = true; + break; + } + } + if (matched) { + break; + } + } + if (matched) { + found.push_back(op); + } + } + return found; +} + class AttentionMaskInput : public ov::pass::MatcherPass { public: OPENVINO_MATCHER_PASS_RTTI("npuw::LLMCompiledModel::AttentionMaskInput"); @@ -123,6 +151,9 @@ class AttentionMaskInput : public ov::pass::MatcherPass { } } } + auto decomposed_cross_attn_nodes = + transform_cross_attn ? find_decomposed_cross_attn_score_nodes(model) + : std::vector>{}; // Self-attention OPENVINO_ASSERT(!self_attn_nodes.empty()); @@ -161,7 +192,7 @@ class AttentionMaskInput : public ov::pass::MatcherPass { if (transform_cross_attn) { // Cross attn - OPENVINO_ASSERT(!cross_attn_nodes.empty()); + OPENVINO_ASSERT(!cross_attn_nodes.empty() || !decomposed_cross_attn_nodes.empty()); // FIXME: Should be taken from topology - don't hardcode!!! auto shape_cst = std::make_shared(ov::element::i64, @@ -191,6 +222,29 @@ class AttentionMaskInput : public ov::pass::MatcherPass { cross_attn_node->input(3).replace_source_output(unsq2->output(0)); } } + for (const auto& qk_score_node : decomposed_cross_attn_nodes) { + if (ov::is_type(qk_score_node)) { + // The original SDPA already had a mask input (or was causal), so + // decomposition produced Add(scores, mask) - mirrors the "cross_attn_node + // already has 4/5 inputs" branch above. + qk_score_node->input(1).replace_source_output(unsq2->output(0)); + } else { + // Mirrors the "cross_attn_node has 3 inputs" branch above: the original + // SDPA had no mask at all, so decomposition produced no Add either - splice + // one in and move the "cross_attention_qk_scaled_scores*" tensor name(s) + // onto it, since that's the value word-level-timestamp extraction expects. + // Snapshot readers before wiring up new_add itself as a reader. + auto readers = qk_score_node->output(0).get_target_inputs(); + auto new_add = std::make_shared(qk_score_node->output(0), unsq2->output(0)); + new_add->set_friendly_name(qk_score_node->get_friendly_name() + "/with_mask"); + auto names = qk_score_node->output(0).get_names(); + qk_score_node->output(0).set_names({}); + new_add->output(0).add_names(names); + for (const auto& reader : readers) { + reader.replace_source_output(new_add->output(0)); + } + } + } } } }; @@ -471,24 +525,81 @@ class WhisperScaledDotProductAttentionDecomposition : public ov::pass::MatcherPa } }; -auto remove_encoder_attn_read_value(const std::shared_ptr& rv_node, - const ov::Output& kv_out, - const ov::Input& sdpa_in) { - // Find Assign node +// A cross-attention KV-cache state (the ReadValue producing the once-computed +// encoder key/value) is consumed differently depending on whether the model +// still has a fused cross-attention SDPA node, or whether cross-attention SDPA +// has already been decomposed (by GenAI, or by +// decompose_scaled_dot_product_attention_for_whisper() below): +// - fused: ReadValue -> [FakeConvert ->] SDPA (key at port 1, value at port 2) +// - decomposed: key feeds the Transpose that builds "kT" for QK^T (and, +// separately, a ShapeOf reading the same state to size that +// transpose - both are redirected, only the Transpose decides the role) +// value feeds the final (softmax @ value) MatMul directly +// Returns the role for a single direct reader of the ReadValue's output, or +// nullopt if that reader alone doesn't tell us (e.g. Assign, ShapeOf). +enum class EncoderKvRole { Key, Value }; + +std::optional classify_encoder_kv_reader(const ov::Input& reader) { + auto* node = reader.get_node(); + if (strstr(node->get_type_name(), "FakeConvert") != nullptr) { + // fp8: ReadValue -> FakeConvert -> {SDPA | Transpose | MatMul}. FakeConvert has a + // single consumer, so its role is exactly the role of that consumer. + auto fc_readers = node->outputs()[0].get_target_inputs(); + OPENVINO_ASSERT(fc_readers.size() == 1); + return classify_encoder_kv_reader(*fc_readers.begin()); + } + if (strstr(node->get_type_name(), "ScaledDotProductAttention") != nullptr) { + return reader.get_index() == 1 ? EncoderKvRole::Key : EncoderKvRole::Value; + } + if (strstr(node->get_type_name(), "Transpose") != nullptr) { + return EncoderKvRole::Key; + } + if (strstr(node->get_type_name(), "MatMul") != nullptr) { + return EncoderKvRole::Value; + } + return std::nullopt; +} + +// Splits an encoder-attn KV-cache ReadValue's readers into its single Assign and the +// (one or more) readers that consume the state's value, and determines whether this +// state is "key" or "value" from among the latter. +struct EncoderKvState { + std::shared_ptr assign_node; + std::vector> value_readers; + EncoderKvRole role; +}; + +EncoderKvState analyze_encoder_kv_read_value(const std::shared_ptr& rv_node) { OPENVINO_ASSERT(rv_node->outputs().size() == 1); - auto rv_out = rv_node->outputs()[0]; - ov::NodeVector rv_readers; - for (const auto& target_in : rv_out.get_target_inputs()) { - rv_readers.push_back(target_in.get_node()->shared_from_this()); + EncoderKvState state; + std::optional role; + for (const auto& reader : rv_node->output(0).get_target_inputs()) { + if (strstr(reader.get_node()->get_type_name(), "Assign") != nullptr) { + OPENVINO_ASSERT(!state.assign_node, "More than one Assign reads an encoder-attn KV-cache state"); + state.assign_node = ov::as_type_ptr(reader.get_node()->shared_from_this()); + continue; + } + state.value_readers.push_back(reader); + if (!role) { + role = classify_encoder_kv_reader(reader); + } } - // Assign and SDPA - OPENVINO_ASSERT(rv_readers.size() == 2); - auto assign_node = (strstr(rv_readers[0]->get_type_name(), "Assign") != nullptr) ? rv_readers[0] : rv_readers[1]; - OPENVINO_ASSERT(strstr(assign_node->get_type_name(), "Assign") != nullptr); - // Redirect KV-cache tensor to SDPA - sdpa_in.replace_source_output(kv_out); - return std::make_pair(std::make_shared(kv_out), - ov::as_type_ptr(assign_node)); + OPENVINO_ASSERT(state.assign_node, "encoder-attn KV-cache state has no Assign"); + OPENVINO_ASSERT(!state.value_readers.empty(), "encoder-attn KV-cache state is never read"); + OPENVINO_ASSERT(role, "Could not classify encoder-attn KV-cache state as key or value"); + state.role = *role; + return state; +} + +auto remove_encoder_attn_read_value(const std::shared_ptr& rv_node, const EncoderKvState& state) { + auto kv_out = rv_node->input_value(0); + // Redirect every consumer of the state directly to its initial value - covers both + // the single fused-SDPA reader and the several decomposed-key readers (Transpose + + // ShapeOf) alike. + for (const auto& reader : state.value_readers) { + reader.replace_source_output(kv_out); + } + return std::make_pair(std::make_shared(kv_out), state.assign_node); } std::string transform_key_value_name(std::string input_string, @@ -511,14 +622,6 @@ void set_name(std::shared_ptr result, const std::string& name) { result->get_output_tensor(0).set_names({name}); } -bool is_fake_cvt_to_key_tensor(const ov::Input& reader) { - auto fc_reader = reader.get_node()->outputs()[0].get_target_inputs(); - // FakeConvert node has only 1 consumer - OPENVINO_ASSERT(fc_reader.size() == 1); - // FakeConvert -> SDPA : 'key' tensor is input with index 1 to SDPA - return fc_reader.begin()->get_index() == 1; -} - void expose_runtime_states_as_outputs(const std::shared_ptr& model) { // Find all ReadValue nodes ov::NodeVector read_value_nodes; @@ -535,35 +638,16 @@ void expose_runtime_states_as_outputs(const std::shared_ptr& model) { // Go through all ReadValue nodes and remove them for (const auto& rv_node : read_value_nodes) { OPENVINO_ASSERT(rv_node->inputs().size() == 1); - OPENVINO_ASSERT(rv_node->outputs().size() == 1); - auto rv_in = rv_node->inputs()[0]; - auto x = rv_in.get_source_output(); - auto rv_out = rv_node->outputs()[0]; - // Gather all nodes that read from ReadValue, there must be SDPA and Assign - auto rv_readers = rv_out.get_target_inputs(); - OPENVINO_ASSERT(rv_readers.size() == 2); - // Input port for SDPA node - for (const auto& reader : rv_readers) { - bool is_fake_cvt = strstr(reader.get_node()->get_type_name(), "FakeConvert") != nullptr; - if (strstr(reader.get_node()->get_type_name(), "ScaledDotProductAttention") != nullptr || is_fake_cvt) { - auto sdpa_in = reader; - - // In case there's additional FakeConvert node(fp8): ReadValue -> FakeConvert -> SDPA - auto is_fc_key_tensor = is_fake_cvt ? is_fake_cvt_to_key_tensor(reader) : false; - - // Remove ReadValue, store new Result and Assign - auto key_or_value = (sdpa_in.get_index() == 1 || is_fc_key_tensor) ? "key" : "value"; - auto [result, assign] = remove_encoder_attn_read_value(rv_node, rv_in.get_source_output(), sdpa_in); - auto normalized_name = - transform_key_value_name(rv_node->inputs()[0].get_source_output().get_node()->get_friendly_name(), - "present", - ".encoder.", - key_or_value); - set_name(result, normalized_name); - results.push_back(result); - assigns.push_back(assign); - } - } + auto state = analyze_encoder_kv_read_value(rv_node); + auto [result, assign] = remove_encoder_attn_read_value(rv_node, state); + auto key_or_value = state.role == EncoderKvRole::Key ? "key" : "value"; + auto normalized_name = transform_key_value_name(rv_node->input_value(0).get_node()->get_friendly_name(), + "present", + ".encoder.", + key_or_value); + set_name(result, normalized_name); + results.push_back(result); + assigns.push_back(assign); } // Add, remove, validate @@ -621,34 +705,25 @@ void expose_runtime_states_as_inputs(const std::shared_ptr& model) { } for (const auto& rv_node : read_value_nodes) { - auto rv_out = rv_node->outputs()[0]; - auto rv_readers = rv_out.get_target_inputs(); - for (auto rv_reader : rv_readers) { - bool is_fake_cvt = strstr(rv_reader.get_node()->get_type_name(), "FakeConvert") != nullptr; - if (strstr(rv_reader.get_node()->get_type_name(), "Assign") != nullptr) { - auto assign_node = ov::as_type_ptr(rv_reader.get_node()->shared_from_this()); - assigns.push_back(assign_node); - } else if (strstr(rv_reader.get_node()->get_type_name(), "ScaledDotProductAttention") != nullptr || - is_fake_cvt) { - auto sdpa_in = rv_reader; - - auto shape = rv_node->get_output_partial_shape(0); - auto new_param = std::make_shared(rv_node->get_output_element_type(0), shape); - - // In case there's additional FakeConvert node(fp8): ReadValue -> FakeConvert -> SDPA - auto is_fc_key_tensor = is_fake_cvt ? is_fake_cvt_to_key_tensor(rv_reader) : false; - - auto key_or_value = (sdpa_in.get_index() == 1 || is_fc_key_tensor) ? "key" : "value"; - auto normalized_name = transform_key_value_name(sdpa_in.get_node()->get_friendly_name(), - "past_key_values", - ".encoder.", - key_or_value); - set_name(new_param, normalized_name); - - params.push_back(new_param); - sdpa_in.replace_source_output(new_param->outputs()[0]); - } + auto state = analyze_encoder_kv_read_value(rv_node); + + auto shape = rv_node->get_output_partial_shape(0); + auto new_param = std::make_shared(rv_node->get_output_element_type(0), shape); + auto key_or_value = state.role == EncoderKvRole::Key ? "key" : "value"; + // Layer index comes from the state's producer (the initial-value subgraph), + // which is unaffected by whether cross-attention SDPA is fused or decomposed - + // unlike the readers, whose names/types differ between the two shapes. + auto normalized_name = transform_key_value_name(rv_node->input_value(0).get_node()->get_friendly_name(), + "past_key_values", + ".encoder.", + key_or_value); + set_name(new_param, normalized_name); + params.push_back(new_param); + + for (const auto& reader : state.value_readers) { + reader.replace_source_output(new_param->output(0)); } + assigns.push_back(state.assign_node); } // Remove sinks and add new params diff --git a/src/plugins/intel_npu/tests/unit/npuw/prepare_whisper_model_test.cpp b/src/plugins/intel_npu/tests/unit/npuw/prepare_whisper_model_test.cpp new file mode 100644 index 00000000000000..ef6e6799f15189 --- /dev/null +++ b/src/plugins/intel_npu/tests/unit/npuw/prepare_whisper_model_test.cpp @@ -0,0 +1,148 @@ +// Copyright (C) 2018-2026 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#include + +#include +#include + +#include "llm_test_helpers.hpp" +#include "openvino/op/ops.hpp" +#include "openvino/op/scaled_dot_product_attention.hpp" +#include "openvino/pass/stateful_to_stateless.hpp" +#include "whisper/prepare_whisper_model.hpp" + +namespace { + +using ov::test::npuw::build_whisper_decoder_test_model; +using ov::test::npuw::WhisperConfig; + +bool has_input_name(const std::shared_ptr& model, const std::string& substr) { + for (const auto& in : model->inputs()) { + for (const auto& name : in.get_names()) { + if (name.find(substr) != std::string::npos) { + return true; + } + } + } + return false; +} + +bool has_output_name(const std::shared_ptr& model, const std::string& substr) { + for (const auto& out : model->outputs()) { + for (const auto& name : out.get_names()) { + if (name.find(substr) != std::string::npos) { + return true; + } + } + } + return false; +} + +// Replicates the query/key/value part of GenAI's WhisperScaledDotProductAttentionDecomposition +// (and NPUW's own copy of it) for a cross-attention SDPA with no mask input - exactly the shape +// build_whisper_decoder_test_model() produces (encoder_attn SDPA has 3 inputs, no explicit mask). +// This is what a model handed to NPUW already looks like once GenAI decomposes cross-attention +// SDPA for NPU too, matching what it already does for CPU/GPU. +std::shared_ptr decompose_one_cross_attn_sdpa(const std::shared_ptr& sdpa) { + using namespace ov::op; + auto query = sdpa->input_value(0); + auto key = sdpa->input_value(1); + auto value = sdpa->input_value(2); + + auto q_shape = std::make_shared(query, ov::element::i32); + auto minus_one = v0::Constant::create(ov::element::i32, ov::Shape{}, {-1}); + auto zero_i = v0::Constant::create(ov::element::i32, ov::Shape{}, {0}); + auto one_i = v0::Constant::create(ov::element::i32, ov::Shape{}, {1}); + auto head_dim = std::make_shared(q_shape, minus_one, zero_i); + auto head_dim_f = std::make_shared(head_dim, query); + auto sqrt_hd = std::make_shared(head_dim_f); + auto one_f = std::make_shared(one_i, query); + auto scale = std::make_shared(one_f, sqrt_hd); + + // Q/K/V are always [batch, heads, seq, head_dim] in this test model - swap the last two axes. + auto perm = v0::Constant::create(ov::element::i32, ov::Shape{4}, {0, 1, 3, 2}); + auto kT = std::make_shared(key, perm); + + auto q_scaled = std::make_shared(query, scale); + auto scaled_atten = std::make_shared(q_scaled, kT); + scaled_atten->output(0).add_names({"cross_attention_qk_scaled_scores"}); + + auto softmax = std::make_shared(scaled_atten, -1); + auto result = std::make_shared(softmax, value); + result->set_friendly_name(sdpa->get_friendly_name()); + return result; +} + +// Decomposes every cross-attention (encoder_attn) SDPA node in-place, simulating a model that +// arrives at NPUW already decomposed by GenAI. +void decompose_cross_attention_sdpa_for_test(const std::shared_ptr& model) { + for (const auto& op : model->get_ordered_ops()) { + auto sdpa = ov::as_type_ptr(op); + if (!sdpa || sdpa->get_friendly_name().find("encoder_attn") == std::string::npos) { + continue; + } + auto new_node = decompose_one_cross_attn_sdpa(sdpa); + ov::replace_node(sdpa, new_node); + } + model->validate_nodes_and_infer_types(); +} + +std::shared_ptr stateless_whisper_decoder_model() { + auto model = build_whisper_decoder_test_model(); + ov::pass::StatefulToStateless().run_on_model(model); + return model->clone(); +} + +class PrepareWhisperModelTest : public ::testing::TestWithParam {}; + +// Runs PrepareWhisperPrefillModel/PrepareWhisperKVCacheModel on both a model with fused +// cross-attention SDPA (today's shape) and one with cross-attention SDPA already decomposed +// (the shape NPUW must also handle once GenAI decomposes for NPU, per CVS-184242). Both must +// succeed and produce the same encoder-attn KV-cache input/output names. +TEST_P(PrepareWhisperModelTest, PrefillPreparationHandlesFusedAndDecomposedCrossAttention) { + const bool decompose = GetParam(); + + auto model = stateless_whisper_decoder_model(); + if (decompose) { + decompose_cross_attention_sdpa_for_test(model); + } + + WhisperConfig cfg; + ASSERT_TRUE(ov::npuw::util::PrepareWhisperPrefillModel(128, + static_cast(cfg.max_source_positions), + false /*decompose_sdpa*/) + .run_on_model(model)); + + EXPECT_TRUE(has_input_name(model, "attention_mask")); + EXPECT_TRUE(has_output_name(model, "present.0.encoder.key")); + EXPECT_TRUE(has_output_name(model, "present.0.encoder.value")); + EXPECT_TRUE(has_output_name(model, "present.1.encoder.key")); + EXPECT_TRUE(has_output_name(model, "present.1.encoder.value")); +} + +TEST_P(PrepareWhisperModelTest, KVCachePreparationHandlesFusedAndDecomposedCrossAttention) { + const bool decompose = GetParam(); + + auto model = stateless_whisper_decoder_model(); + if (decompose) { + decompose_cross_attention_sdpa_for_test(model); + } + + ASSERT_TRUE(ov::npuw::util::PrepareWhisperKVCacheModel().run_on_model(model)); + + EXPECT_TRUE(has_input_name(model, "past_key_values.0.encoder.key")); + EXPECT_TRUE(has_input_name(model, "past_key_values.0.encoder.value")); + EXPECT_TRUE(has_input_name(model, "past_key_values.1.encoder.key")); + EXPECT_TRUE(has_input_name(model, "past_key_values.1.encoder.value")); +} + +INSTANTIATE_TEST_SUITE_P(CrossAttentionShape, + PrepareWhisperModelTest, + ::testing::Bool(), + [](const ::testing::TestParamInfo& info) { + return info.param ? "Decomposed" : "Fused"; + }); + +} // namespace