From 596c1309f820ef5917a78e89fcb0edb99916e071 Mon Sep 17 00:00:00 2001 From: Alexis Menard Date: Thu, 13 Aug 2026 15:49:18 -0400 Subject: [PATCH] [NPUW] validate pyramid attention port indices on deserialize ### Tickets: - [CVS-192805](https://jira.devtools.intel.com/browse/CVS-192805) ### Details: - Validate pyramid-attention port indices after deserializing compiled models and before using them at runtime. - Reject malformed serialized pyramid state instead of allowing out-of-bounds access into compiled model inputs. - The auto mask_iport assignment now uses `pyramid->_compiled_models[pyramid_id]->inputs().at(dyn_mask_idx);` to ensure invalid serialized metadata fails cleanly with `ov::Exception` instead of indexing past the end of the model input vector. - Add regression coverage for valid, invalid, and malformed serialized pyramid-attention states to lock in the fix for the issue described in CVS-192805. - Preserve the runtime guard in the import/deserialization path where the corrupted pyramid state is rehydrated and attached to compiled models. ### AI Assistance: - AI assistance used: yes - AI found the issue and helped generate the fix and regression coverage. Human validated the regression coverage. --- .../src/plugin/npuw/attn/attn_subgraph.cpp | 3 +- .../src/plugin/npuw/pyramid_attention.cpp | 63 +++++ .../src/plugin/npuw/pyramid_attention.hpp | 7 + .../unit/npuw/pyramid_attention_test.cpp | 240 ++++++++++++++++++ 4 files changed, 312 insertions(+), 1 deletion(-) diff --git a/src/plugins/intel_npu/src/plugin/npuw/attn/attn_subgraph.cpp b/src/plugins/intel_npu/src/plugin/npuw/attn/attn_subgraph.cpp index 6f32a58888977e..05343347581988 100644 --- a/src/plugins/intel_npu/src/plugin/npuw/attn/attn_subgraph.cpp +++ b/src/plugins/intel_npu/src/plugin/npuw/attn/attn_subgraph.cpp @@ -792,7 +792,7 @@ ov::npuw::v1::subgraphs::RuntimeBehaviorFactory make_runtime_factory() { const auto pyramid_id = state.pyramid_selector->pyramid_id(); const std::size_t dyn_mask_idx = pyramid->mask_idx_at(pyramid_id); const std::size_t dyn_query_size = pyramid->query_size_at(pyramid_id); - auto mask_iport = pyramid->_compiled_models[pyramid_id]->inputs()[dyn_mask_idx]; + auto mask_iport = pyramid->_compiled_models[pyramid_id]->inputs().at(dyn_mask_idx); const auto& graph_mask = io.inputs.at(dyn_mask_idx); const auto this_case = state.pyramid_selector->this_case(); const auto present_len = dyn_query_size; @@ -1218,6 +1218,7 @@ void serialize_compiled_state(v1::subgraphs::Context& context, mutable_pyramid->_compiled_models[num_models - 1] = submodel_ctx->compiled_model; LOG_DEBUG("Reused compiled_model for the last pyramid attention model"); } + mutable_pyramid->validate_port_indices(); } } diff --git a/src/plugins/intel_npu/src/plugin/npuw/pyramid_attention.cpp b/src/plugins/intel_npu/src/plugin/npuw/pyramid_attention.cpp index fbf6dfd1bdfe72..0137d71a6186e6 100644 --- a/src/plugins/intel_npu/src/plugin/npuw/pyramid_attention.cpp +++ b/src/plugins/intel_npu/src/plugin/npuw/pyramid_attention.cpp @@ -814,6 +814,69 @@ void PyramidAttentionContiguous::collect_strided_input_names(const ov::Model& mo } } +void PyramidAttentionContiguous::validate_port_indices() const { + if (_attention_infos.size() != _compiled_models.size()) { + OPENVINO_THROW("NPU NPUW: pyramid attention info count (", + _attention_infos.size(), + ") does not match compiled model count (", + _compiled_models.size(), + ")"); + } + for (size_t i = 0; i < _compiled_models.size(); ++i) { + if (!_compiled_models[i]) { + continue; + } + const auto inputs_size = _compiled_models[i]->inputs().size(); + const auto& info = _attention_infos[i]; + if (info.mask_idx >= inputs_size) { + OPENVINO_THROW("NPU NPUW: pyramid attention mask_idx (", + info.mask_idx, + ") out of bounds for model ", + i, + " with ", + inputs_size, + " inputs"); + } + for (const auto& param : info.params) { + if (param.idx >= inputs_size) { + OPENVINO_THROW("NPU NPUW: pyramid attention param idx (", + param.idx, + ") out of bounds for model ", + i, + " with ", + inputs_size, + " inputs"); + } + } + } +} + +void PyramidAttentionBlock::validate_port_indices() const { + if (_attention_infos.size() != _compiled_models.size()) { + OPENVINO_THROW("NPU NPUW: pyramid attention info count (", + _attention_infos.size(), + ") does not match compiled model count (", + _compiled_models.size(), + ")"); + } + for (size_t i = 0; i < _compiled_models.size(); ++i) { + if (!_compiled_models[i]) { + continue; + } + const auto inputs_size = _compiled_models[i]->inputs().size(); + const auto& info = _attention_infos[i]; + if (info.mask_idx >= inputs_size) { + OPENVINO_THROW("NPU NPUW: pyramid attention mask_idx (", + info.mask_idx, + ") out of bounds for model ", + i, + " with ", + inputs_size, + " inputs"); + } + } +} + // ── PyramidAttention::make() static factory ─────────────────────────────────────── std::shared_ptr PyramidAttention::make(const function::PyramidAttention& func_pyramid) { diff --git a/src/plugins/intel_npu/src/plugin/npuw/pyramid_attention.hpp b/src/plugins/intel_npu/src/plugin/npuw/pyramid_attention.hpp index 00722502b519c2..6ef1b913ad41a6 100644 --- a/src/plugins/intel_npu/src/plugin/npuw/pyramid_attention.hpp +++ b/src/plugins/intel_npu/src/plugin/npuw/pyramid_attention.hpp @@ -220,6 +220,11 @@ struct PyramidAttention { // No-op in block mode (block ports are bound directly, not via strided views). virtual void collect_strided_input_names(const ov::Model& model, std::string& out) const = 0; + // Validates that all port indices in _attention_infos are within bounds of the + // corresponding compiled model's inputs(). Throws ov::Exception on violation. + // Must be called after _compiled_models is fully populated (import path). + virtual void validate_port_indices() const = 0; + // Non-virtual shared methods void set_compiled_models(std::vector>&& compiled_models); @@ -264,6 +269,7 @@ struct PyramidAttentionContiguous final : PyramidAttention { } std::optional kv_param_dim(size_t pyramid_id, size_t input_idx) const override; void collect_strided_input_names(const ov::Model& model, std::string& out) const override; + void validate_port_indices() const override; }; // Concrete subclass for block-split KV cache mode (after SplitKVCacheIntoBlocks). @@ -307,6 +313,7 @@ struct PyramidAttentionBlock final : PyramidAttention { return std::nullopt; } void collect_strided_input_names(const ov::Model&, std::string&) const override {} // no-op + void validate_port_indices() const override; }; } // namespace compiled diff --git a/src/plugins/intel_npu/tests/unit/npuw/pyramid_attention_test.cpp b/src/plugins/intel_npu/tests/unit/npuw/pyramid_attention_test.cpp index bd49d48c8d82dc..b0302fa63a9cbd 100644 --- a/src/plugins/intel_npu/tests/unit/npuw/pyramid_attention_test.cpp +++ b/src/plugins/intel_npu/tests/unit/npuw/pyramid_attention_test.cpp @@ -10,6 +10,7 @@ #include #include +#include "attn/attn_subgraph.hpp" #include "openvino/op/add.hpp" #include "openvino/op/concat.hpp" #include "openvino/op/constant.hpp" @@ -21,6 +22,7 @@ #include "npuw_transformations/convert_kvcache_to_precision.hpp" #include "npuw_transformations/split_kvcache_into_blocks.hpp" #include "pyramid_attention.hpp" +#include "serialization.hpp" #include "util.hpp" namespace { @@ -823,4 +825,242 @@ TEST(PyramidAttentionTest, ProcessBlockModePyramidModelsProducesGrowingContextLe } } +// --- Security regression tests: CWE-125 OOB read via attacker-controlled mask_idx --- +// Regression for: deserialized pyramid attention port indices must be validated against +// the compiled model's inputs().size() before use as vector subscripts. +// +// Each test below goes through the same production code path used in attn_subgraph.cpp: +// export: orc::serialize(Stream::writer, PyramidAttentionContiguous) +// import: make_pyramid_from_stream(Stream::reader, tag) <- same call as attn_subgraph.cpp:1185 +// attach models, then call validate_port_indices() <- same call added to attn_subgraph.cpp + +namespace { + +// Minimal IPlugin stub needed to construct an ICompiledModel. +class NullPluginStub final : public ov::IPlugin { +public: + std::shared_ptr compile_model(const std::shared_ptr&, + const ov::AnyMap&) const override { + return {}; + } + std::shared_ptr compile_model(const std::shared_ptr&, + const ov::AnyMap&, + const ov::SoPtr&) const override { + return {}; + } + std::shared_ptr import_model(std::istream&, const ov::AnyMap&) const override { + return {}; + } + std::shared_ptr import_model(std::istream&, + const ov::SoPtr&, + const ov::AnyMap&) const override { + return {}; + } + std::shared_ptr import_model(const ov::Tensor&, const ov::AnyMap&) const override { + return {}; + } + std::shared_ptr import_model(const ov::Tensor&, + const ov::SoPtr&, + const ov::AnyMap&) const override { + return {}; + } + ov::SupportedOpsMap query_model(const std::shared_ptr&, const ov::AnyMap&) const override { + return {}; + } + void set_property(const ov::AnyMap&) override {} + ov::Any get_property(const std::string&, const ov::AnyMap&) const override { + return {}; + } + ov::SoPtr create_context(const ov::AnyMap&) const override { + return {}; + } + ov::SoPtr get_default_context(const ov::AnyMap&) const override { + return {}; + } +}; + +// Minimal ICompiledModel stub that exposes the inputs of the wrapped ov::Model. +class StubCompiledModel final : public ov::ICompiledModel { +public: + StubCompiledModel(const std::shared_ptr& model, const std::shared_ptr& plugin) + : ov::ICompiledModel(model, plugin) {} + + 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 {}; + } + std::shared_ptr create_sync_infer_request() const override { + return nullptr; + } +}; + +// Build a simple model with n_inputs parameters and return a compiled stub. +ov::SoPtr make_stub_model(size_t n_inputs, const std::shared_ptr& plugin) { + ov::ParameterVector params; + for (size_t i = 0; i < n_inputs; ++i) { + auto p = std::make_shared(ov::element::f32, ov::Shape{1}); + p->set_friendly_name("input_" + std::to_string(i)); + p->output(0).get_tensor().set_names({"input_" + std::to_string(i)}); + params.push_back(p); + } + auto result = std::make_shared(params[0]); + auto model = std::make_shared(ov::ResultVector{result}, params); + return {std::make_shared(model, plugin), {}}; +} + +template +std::shared_ptr import_pyramid_from_stream(Pyramid& src, uint8_t tag) { + std::stringstream ss(std::ios::in | std::ios::out | std::ios::binary); + { + auto wstream = ov::npuw::orc::Stream::writer(ss); + wstream & tag; + ov::npuw::orc::serialize(wstream, src); + } + ss.seekg(0); + auto rstream = ov::npuw::orc::Stream::reader(ss); + uint8_t read_tag = 0u; + rstream & read_tag; + return ov::npuw::orc::make_pyramid_from_stream(rstream, read_tag); +} + +void attach_compiled_models(std::shared_ptr pyramid, + std::initializer_list model_input_sizes, + const std::shared_ptr& plugin) { + pyramid->_compiled_models.clear(); + for (const auto input_size : model_input_sizes) { + pyramid->_compiled_models.push_back(make_stub_model(input_size, plugin)); + } +} + +template +void expect_invalid_port_indices_rejected(Pyramid& src, + uint8_t tag, + std::initializer_list model_input_sizes, + const std::string& case_name) { + SCOPED_TRACE(case_name); + auto pyramid = import_pyramid_from_stream(src, tag); + ASSERT_NE(pyramid, nullptr); + + auto plugin = std::make_shared(); + attach_compiled_models(pyramid, model_input_sizes, plugin); + EXPECT_THROW(pyramid->validate_port_indices(), ov::Exception); +} + +} // namespace + +TEST(PyramidAttentionTest, ValidPortIndicesPassValidation) { + ov::npuw::compiled::PyramidAttentionContiguous src; + src.query_size = 1; + src.full_context_size = 64; + src._context_lengths = {64}; + ov::npuw::compiled::PyramidAttentionContiguousInfo info; + info.mask_idx = 2; + info.params = {{1, 0}}; + src._attention_infos = {info}; + + auto pyramid = import_pyramid_from_stream(src, 0u); + ASSERT_NE(pyramid, nullptr); + + auto plugin = std::make_shared(); + pyramid->_compiled_models = {make_stub_model(5, plugin)}; + + EXPECT_NO_THROW(pyramid->validate_port_indices()); +} + +TEST(PyramidAttentionTest, MalformedSerializedPyramidStateIsRejectedOnDeserialize) { + using ContigInfo = ov::npuw::compiled::PyramidAttentionContiguousInfo; + + auto plugin = std::make_shared(); + auto stub_model = make_stub_model(3, plugin); + + auto src = std::make_shared(); + src->query_size = 1; + src->full_context_size = 64; + src->_context_lengths = {64}; + ContigInfo info; + info.mask_idx = 0xFF; + src->_attention_infos = {info}; + src->_compiled_models = {stub_model}; + + ov::npuw::v1::subgraphs::Context serialized_context; + ov::npuw::attn::put_compiled_pyramid(serialized_context, src); + + std::stringstream ss(std::ios::in | std::ios::out | std::ios::binary); + { + auto stream = ov::npuw::s11n::Stream::writer(ss); + ov::npuw::attn::serialize_compiled_state(serialized_context, stream, nullptr); + } + + ss.seekg(0); + ov::npuw::v1::subgraphs::Context restored_context; + auto imported_model = make_stub_model(3, plugin); + ov::npuw::s11n::SubmodelDeserializeCtx submodel_ctx(plugin, "CPU", imported_model); + + { + auto stream = ov::npuw::s11n::Stream::reader(ss); + EXPECT_THROW(ov::npuw::attn::serialize_compiled_state(restored_context, stream, &submodel_ctx), ov::Exception) + << "Malformed serialized pyramid attention metadata must be rejected during deserialization"; + } +} + +TEST(PyramidAttentionTest, InvalidPortIndicesAreRejected) { + using ContigInfo = ov::npuw::compiled::PyramidAttentionContiguousInfo; + using BlockInfo = ov::npuw::compiled::PyramidAttentionBlockInfo; + + { + ov::npuw::compiled::PyramidAttentionContiguous src; + src.query_size = 1; + src.full_context_size = 64; + src._context_lengths = {64}; + ContigInfo info; + info.mask_idx = 0xFF; + src._attention_infos = {info}; + expect_invalid_port_indices_rejected(src, 0u, {3}, "mask_idx out of range"); + } + + { + ov::npuw::compiled::PyramidAttentionContiguous src; + src.query_size = 1; + src.full_context_size = 64; + src._context_lengths = {64}; + ContigInfo info; + info.mask_idx = 0; + info.params = {{0xFFFFFFFF, 0}}; + src._attention_infos = {info}; + expect_invalid_port_indices_rejected(src, 0u, {3}, "param idx out of range"); + } + + { + ov::npuw::compiled::PyramidAttentionBlock src; + src.query_size = 1; + src.full_context_size = 64; + src._context_lengths = {64}; + BlockInfo info; + info.mask_idx = 0xFF; + src._attention_infos = {info}; + expect_invalid_port_indices_rejected(src, 1u, {4}, "block mask idx out of range"); + } + + { + ov::npuw::compiled::PyramidAttentionContiguous src; + src.query_size = 1; + src.full_context_size = 64; + src._context_lengths = {64}; + ContigInfo info; + info.mask_idx = 0; + src._attention_infos = {info}; + // Model count mismatches the serialized attention info count. + auto pyramid = import_pyramid_from_stream(src, 0u); + ASSERT_NE(pyramid, nullptr); + + auto plugin = std::make_shared(); + pyramid->_compiled_models = {make_stub_model(3, plugin), make_stub_model(3, plugin)}; + EXPECT_THROW(pyramid->validate_port_indices(), ov::Exception); + } +} + } // namespace