[NPUW] validate pyramid attention port indices on deserialize

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

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> PyramidAttention::make(const function::PyramidAttention& func_pyramid) {

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<ov::SoPtr<ov::ICompiledModel>>&& compiled_models);
 
@@ -264,6 +269,7 @@ struct PyramidAttentionContiguous final : PyramidAttention {
     }
     std::optional<std::size_t> 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

src/plugins/intel_npu/tests/unit/npuw/pyramid_attention_test.cpp

@@ -10,6 +10,7 @@
 #include <variant>
 #include <vector>
 
+#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<ov::ICompiledModel> compile_model(const std::shared_ptr<const ov::Model>&,
+                                                      const ov::AnyMap&) const override {
+        return {};
+    }
+    std::shared_ptr<ov::ICompiledModel> compile_model(const std::shared_ptr<const ov::Model>&,
+                                                      const ov::AnyMap&,
+                                                      const ov::SoPtr<ov::IRemoteContext>&) const override {
+        return {};
+    }
+    std::shared_ptr<ov::ICompiledModel> import_model(std::istream&, const ov::AnyMap&) const override {
+        return {};
+    }
+    std::shared_ptr<ov::ICompiledModel> import_model(std::istream&,
+                                                     const ov::SoPtr<ov::IRemoteContext>&,
+                                                     const ov::AnyMap&) const override {
+        return {};
+    }
+    std::shared_ptr<ov::ICompiledModel> import_model(const ov::Tensor&, const ov::AnyMap&) const override {
+        return {};
+    }
+    std::shared_ptr<ov::ICompiledModel> import_model(const ov::Tensor&,
+                                                     const ov::SoPtr<ov::IRemoteContext>&,
+                                                     const ov::AnyMap&) const override {
+        return {};
+    }
+    ov::SupportedOpsMap query_model(const std::shared_ptr<const ov::Model>&, 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<ov::IRemoteContext> create_context(const ov::AnyMap&) const override {
+        return {};
+    }
+    ov::SoPtr<ov::IRemoteContext> 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<ov::Model>& model, const std::shared_ptr<const ov::IPlugin>& plugin)
+        : ov::ICompiledModel(model, plugin) {}
+
+    void export_model(std::ostream&) const override {}
+    std::shared_ptr<const ov::Model> 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<ov::ISyncInferRequest> create_sync_infer_request() const override {
+        return nullptr;
+    }
+};
+
+// Build a simple model with n_inputs parameters and return a compiled stub.
+ov::SoPtr<ov::ICompiledModel> make_stub_model(size_t n_inputs, const std::shared_ptr<const ov::IPlugin>& plugin) {
+    ov::ParameterVector params;
+    for (size_t i = 0; i < n_inputs; ++i) {
+        auto p = std::make_shared<ov::op::v0::Parameter>(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<ov::op::v0::Result>(params[0]);
+    auto model = std::make_shared<ov::Model>(ov::ResultVector{result}, params);
+    return {std::make_shared<StubCompiledModel>(model, plugin), {}};
+}
+
+template <typename Pyramid>
+std::shared_ptr<ov::npuw::compiled::PyramidAttention> 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<ov::npuw::compiled::PyramidAttention> pyramid,
+                            std::initializer_list<size_t> model_input_sizes,
+                            const std::shared_ptr<const ov::IPlugin>& 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 <typename Pyramid>
+void expect_invalid_port_indices_rejected(Pyramid& src,
+                                          uint8_t tag,
+                                          std::initializer_list<size_t> 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<NullPluginStub>();
+    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<NullPluginStub>();
+    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<NullPluginStub>();
+    auto stub_model = make_stub_model(3, plugin);
+
+    auto src = std::make_shared<ov::npuw::compiled::PyramidAttentionContiguous>();
+    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<NullPluginStub>();
+        pyramid->_compiled_models = {make_stub_model(3, plugin), make_stub_model(3, plugin)};
+        EXPECT_THROW(pyramid->validate_port_indices(), ov::Exception);
+    }
+}
+
 }  // namespace