[NPU] Weight sharing with cache: Part 1 make compiled model weights WeightContext friendly

src/plugins/intel_npu/src/compiler_adapter/src/weightless_graph.cpp

@@ -459,6 +459,10 @@ WeightlessGraph::InputData WeightlessGraph::allocate_inputs(
 
     // Due to the large number of init inputs, allocating a single buffer for all of them is more efficient. "View
     // tensors" are used for separating them.
+    // TODO in case of cross-context weight sharing the current implementation has been allocating a single buffer for several weights coul be shared
+    // In this case we don't need to allocate the signle buffer for such weight, we can reuse it instead.
+    // But we following the current logic, we still have to allocate a single buffer for non-shared weights.
+    // Implication: we need to return multiple ZeroTensors instead of a single one.
     const std::shared_ptr<ZeroTensor> initInputsAllocatedTensor =
         std::make_shared<ZeroTensor>(_zeroInitStruct, ov::element::Type_t::u8, ov::Shape({initInputsByteSize}), true);
 

src/plugins/intel_npu/src/plugin/include/shared_weights_assigner.hpp

@@ -0,0 +1,82 @@
+// Copyright (C) 2018-2026 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include <functional>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "openvino/runtime/aligned_buffer.hpp"
+
+namespace ov {
+class Model;
+namespace op {
+namespace v0 {
+class Constant;
+}  // namespace v0
+}  // namespace op
+}  // namespace ov
+
+namespace ov {
+namespace intel_npu {
+
+class SharedWeightsAssigner {
+public:
+    using SharedConstant = std::shared_ptr<ov::op::v0::Constant>;
+    using PartitionedConstants = std::vector<std::vector<SharedConstant>>;
+    using SharedSourceAndConstants = std::pair<std::shared_ptr<ov::AlignedBuffer>, std::vector<SharedConstant>>;
+    using SharedSourcesWithConstants = std::vector<SharedSourceAndConstants>;
+
+    struct Options {
+        std::vector<std::string> shared_device_contexts;
+        size_t single_weight_shared_source_size_max = 0;
+        bool preserve_weightless_cache_attr = true;
+        std::function<size_t()> source_id_generator;
+    };
+
+    struct Statistic {
+        std::vector<size_t> partition_constant_counts;
+        size_t collected_constants_count = 0;
+        size_t total_shared_constant_bytes = 0;
+        size_t total_non_shared_constant_bytes_released = 0;
+
+        std::string to_string() const;
+    };
+
+    struct CollectResult {
+        Statistic statistic;
+        PartitionedConstants partitioned_constants;
+    };
+
+    explicit SharedWeightsAssigner(Options options);
+    CollectResult collect_and_partition(const std::shared_ptr<ov::Model>& model);
+    SharedSourcesWithConstants mutate_model_with_constant_sharing(PartitionedConstants&& partitioned_constants);
+
+private:
+    bool constant_can_be_shared(const ov::op::v0::Constant& constant) const;
+
+    std::vector<SharedConstant> collect_weights_to_share(const std::shared_ptr<ov::Model>& model) const;
+
+    PartitionedConstants partition_constants_by_size(std::vector<SharedConstant>&& constants) const;
+
+    SharedSourcesWithConstants make_constant_shareable(
+        PartitionedConstants&& partitioned_constants) const;
+
+    std::shared_ptr<ov::AlignedBuffer> make_shared_source(const std::vector<SharedConstant>& partition) const;
+
+    size_t get_constant_aligned_size(const ov::op::v0::Constant& constant) const;
+
+    static size_t align_bytes(size_t bytes, size_t alignment);
+
+    std::vector<std::string> m_shared_device_contexts;
+    std::function<size_t()> m_source_id_generator;
+    bool m_preserve_weightless_cache_attr = true;
+    size_t m_min_relocate_bytes = 0;
+    size_t m_single_weight_shared_source_size_max = 0;
+};
+
+}  // namespace intel_npu
+}  // namespace ov

src/plugins/intel_npu/src/plugin/npuw/llm_compiled_model.cpp

@@ -1,6 +1,8 @@
 // Copyright (C) 2018-2026 Intel Corporation
 // SPDX-License-Identifier: Apache-2.0
 //
+
+#include <queue>
 #include "llm_compiled_model.hpp"
 
 #include "embedding/embedding_infer_request.hpp"
@@ -42,12 +44,15 @@
 #include "openvino/pass/pattern/op/wrap_type.hpp"
 #include "openvino/pass/stateful_to_stateless.hpp"
 #include "openvino/pass/validate.hpp"
+#include "openvino/runtime/device_id_parser.hpp"
 #include "openvino/runtime/iasync_infer_request.hpp"
+#include "openvino/runtime/internal_properties.hpp"
 #include "openvino/runtime/properties.hpp"
 #include "partitioning/patterns/fold_const.hpp"
 #include "partitioning/patterns/moe.hpp"
 #include "partitioning/patterns/pre_compute.hpp"
 #include "partitioning/patterns/sdpa.hpp"
+#include "shared_weights_assigner.hpp"
 #include "serialization.hpp"
 #include "transformations/convert_precision.hpp"
 #include "util.hpp"
@@ -684,6 +689,47 @@ void ov::npuw::LLMCompiledModel::compile_generate_model_variants(
     }
 }
 
+void ov::npuw::LLMCompiledModel::assign_shared_weight_to_model_if_possible(const std::shared_ptr<ov::Model> model, const std::shared_ptr<const ov::IPlugin>& plugin,
+const ov::AnyMap& properties) {
+    constexpr size_t single_weight_shared_source_size_max = static_cast<size_t>(2ULL * 1024 * 1024 * 1024);
+
+    NPUW_ASSERT(model && "Model for assigning shared weights must not be null");
+    NPUW_ASSERT(plugin && "Plugin for assigning shared weights must not be null");
+    auto shared_weight_property_it = properties.find("SHARED_WEIGHTS");
+    if (shared_weight_property_it == properties.end()) {
+        return;
+    }
+
+    auto shared_device_contexts =
+        ov::DeviceIDParser::get_hetero_devices(shared_weight_property_it->second.as<std::string>());
+    ::ov::intel_npu::SharedWeightsAssigner::Options shared_weights_assigner_options;
+    shared_weights_assigner_options.shared_device_contexts = std::move(shared_device_contexts);
+    shared_weights_assigner_options.single_weight_shared_source_size_max = single_weight_shared_source_size_max;
+    shared_weights_assigner_options.preserve_weightless_cache_attr = (std::getenv("NO_WEIGHTLESS_ATTR") == nullptr);
+    ::ov::intel_npu::SharedWeightsAssigner shared_weights_assigner(std::move(shared_weights_assigner_options));
+    auto collect_result = shared_weights_assigner.collect_and_partition(model);
+
+    LOG_INFO("[NPUW] SHARED_WEIGHTS: " << collect_result.statistic.to_string());
+    for (size_t i = 0; i < collect_result.statistic.partition_constant_counts.size(); ++i) {
+        LOG_INFO("[NPUW] SHARED_WEIGHTS: partition " << i + 1 << "/"
+                 << collect_result.statistic.partition_constant_counts.size()
+                 << ", constants count: " << collect_result.statistic.partition_constant_counts[i]);
+    }
+
+    auto shared_sources_with_constants =
+        shared_weights_assigner.mutate_model_with_constant_sharing(std::move(collect_result.partitioned_constants));
+
+    m_shared_weight_sources.clear();
+    for (const auto& [shared_source, constants] : shared_sources_with_constants) {
+        LOG_INFO("[NPUW] SHARED_WEIGHTS: allocated shared source buffer: source_id: " << shared_source->get_descriptor()->get_id() 
+                 << ", ptr: " << static_cast<void*>(shared_source->get_ptr<char>())
+                 << ", size: " << shared_source->size()
+                 << ", holds shared weight count: " << constants.size());
+        // Keep source buffers alive for the lifetime of this compiled model.
+        m_shared_weight_sources.push_back(shared_source);
+    }
+}
+
 ov::npuw::LLMCompiledModel::LLMCompiledModel(const std::shared_ptr<ov::Model>& model,
                                              const std::shared_ptr<const ov::IPlugin>& plugin,
                                              const ov::AnyMap& properties,
@@ -820,6 +866,9 @@ ov::npuw::LLMCompiledModel::LLMCompiledModel(const std::shared_ptr<ov::Model>& m
         LOG_INFO("Continuous prefill is enabled");
     }
 
+    LOG_DEBUG("Assigning shared weights to model if possible.");
+    assign_shared_weight_to_model_if_possible(model, plugin, properties);
+
     const uint32_t batch_dim = m_cfg.get<::intel_npu::NPUW_LLM_BATCH_DIM>();
     const uint32_t seq_len_dim = m_cfg.get<::intel_npu::NPUW_LLM_SEQ_LEN_DIM>();
     KVAxesPosition axes{batch_dim, seq_len_dim};

src/plugins/intel_npu/src/plugin/npuw/llm_compiled_model.hpp

@@ -8,6 +8,7 @@
 
 #include "compiled_model.hpp"
 #include "npuw_transformations/kv_axes_position.hpp"
+#include "openvino/core/weight_sharing_util.hpp"
 
 namespace ov {
 namespace test {
@@ -168,7 +169,14 @@ class LLMCompiledModel : public ov::npuw::ICompiledModel {
                                          const std::shared_ptr<const ov::IPlugin>& plugin,
                                          const ov::AnyMap& generate_config);
 
+    // Relocate page-aligned constants into shared cross-device memory if SHARED_WEIGHTS is set.
+    void assign_shared_weight_to_model_if_possible(const std::shared_ptr<ov::Model> model,
+                                                   const std::shared_ptr<const ov::IPlugin>& plugin,
+                                                   const ov::AnyMap& properties);
+
     bool m_is_eagle = false;
+
+    std::vector<std::shared_ptr<AlignedBuffer>> m_shared_weight_sources;
 };
 
 }  // namespace npuw

src/plugins/intel_npu/src/plugin/sources.cmake

@@ -14,6 +14,7 @@ set(SOURCES
     ${CMAKE_CURRENT_SOURCE_DIR}/include/plugin_property_manager.hpp
     ${CMAKE_CURRENT_SOURCE_DIR}/include/property_registration.hpp
     ${CMAKE_CURRENT_SOURCE_DIR}/include/remote_context.hpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/include/shared_weights_assigner.hpp
     ${CMAKE_CURRENT_SOURCE_DIR}/include/transformations.hpp
     ${CMAKE_CURRENT_SOURCE_DIR}/src/async_infer_request.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/src/backends_registry.cpp
@@ -26,6 +27,7 @@ set(SOURCES
     ${CMAKE_CURRENT_SOURCE_DIR}/src/plugin.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/src/plugin_property_manager.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/src/remote_context.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/src/shared_weights_assigner.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/src/transformations.cpp
 )