[NPUW]Annotate each Add in SDPA to make per-ATTN-subgraph mask-skipping decisions

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

@@ -11,6 +11,7 @@
 
 #include "intel_npu/ops/flash_attention_tile.hpp"
 #include "logging.hpp"
+#include "npuw_transformations/detect_causal_mask.hpp"
 #include "openvino/core/validation_util.hpp"
 #include "openvino/op/ops.hpp"
 #include "openvino/openvino.hpp"
@@ -1106,6 +1107,32 @@ std::optional<HostFlashAttention> HostFlashAttention::from(const std::shared_ptr
     LOG_INFO("Creating HFA tile models: tile_size=" << query_size << ", v_transposed=" << v_transposed
                                                     << ", block_kv=" << block_kv_dtype
                                                     << ", present_kv=" << present_kv_dtype << ", q=" << q_dtype);
+
+    // Per-SDPA mask-skipping override
+    // AnnotatePerSDPAMaskType may have annotated this subgraph's Add(QK, mask) node
+    // with its individual mask type. For mixed SWA + global-attention models
+    // (e.g. Gemma-4 E2B/E4B), each ATTN subgraph decides independently:
+    //
+    //   Causal: force enable_mask_skipping = true
+    //   SlidingWindow / no annotation: keep the global enable_mask_skipping unchanged
+    //
+    // SlidingWindow is not forced to false because the global flag already handles the
+    // window_size >= max_prompt_len case (wide SWA that covers the full context).
+    bool local_enable_mask_skipping = enable_mask_skipping;
+    if (pattern_nodes.add_node) {
+        const auto& rt_info = pattern_nodes.add_node->get_rt_info();
+        const auto it = rt_info.find(ov::npuw::NPUW_SDPA_MASK_TYPE_RT_KEY);
+        if (it != rt_info.end()) {
+            const auto per_sdpa_mask_type = static_cast<ov::npuw::MaskInfo::MaskType>(it->second.as<int>());
+            if (per_sdpa_mask_type == ov::npuw::MaskInfo::MaskType::Causal) {
+                local_enable_mask_skipping = true;
+                LOG_DEBUG("Per-SDPA mask annotation: Causal → mask skipping ENABLED for this ATTN subgraph");
+            } else {
+                LOG_DEBUG("Per-SDPA mask annotation: SlidingWindow/Unknown → use global mask skipping setting ("
+                          << (enable_mask_skipping ? "YES" : "NO") << ") for this ATTN subgraph");
+            }
+        }
+    }
     auto tile_model = create_hfa_tile_model(q_shape_static,
                                             block_kv_dtype,  // state_dtype
                                             block_kv_dtype,  // kv_tile_dtype (past blocks)
@@ -1115,7 +1142,7 @@ std::optional<HostFlashAttention> HostFlashAttention::from(const std::shared_ptr
                                             kv_num_heads,
                                             false,
                                             fused_flash_attention,
-                                            enable_mask_skipping,
+                                            local_enable_mask_skipping,
                                             v_transposed);
     if (!tile_model) {
         LOG_WARN("Failed to create HFA tile model");
@@ -1131,7 +1158,7 @@ std::optional<HostFlashAttention> HostFlashAttention::from(const std::shared_ptr
                                                   kv_num_heads,
                                                   true,
                                                   fused_flash_attention,
-                                                  enable_mask_skipping,
+                                                  local_enable_mask_skipping,
                                                   v_transposed,
                                                   output_dtype);
     if (!final_tile_model) {

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

@@ -1051,6 +1051,17 @@ ov::npuw::LLMCompiledModel::LLMCompiledModel(const std::shared_ptr<ov::Model>& m
     } else {
         LOG_DEBUG("Check and apply opt layout --- SKIPPED");
     }
+
+    // Annotate each Add(QK, mask) node with its per-SDPA mask type via rt_info.
+    // Must run AFTER OptimizeValueTensors because ScaledDotProductAttentionDecomposition
+    // inside it is what creates the Add(QK, mask) nodes from SDPA ops.
+    // For mixed SWA + global-attention models (e.g. Gemma-4 E2B/E4B), HFA uses these
+    // annotations to make per-ATTN-subgraph mask-skipping decisions.
+    ov::npuw::AnnotatePerSDPAMaskType().run_on_model(prefill_model);
+    for (auto& model_variant : generate_model_variants) {
+        ov::npuw::AnnotatePerSDPAMaskType().run_on_model(model_variant);
+    }
+
     if (!m_is_embedding) {
         if (!m_use_chunk_prefill) {
             LOG_DEBUG("Removing EmptyKVInputs");

src/plugins/intel_npu/src/plugin/npuw/npuw_transformations/detect_causal_mask.cpp

@@ -5,7 +5,10 @@
 #include "detect_causal_mask.hpp"
 
 #include <cstdlib>
+#include <queue>
+#include <unordered_set>
 
+#include "../util.hpp"
 #include "openvino/op/ops.hpp"
 #include "openvino/op/scaled_dot_product_attention.hpp"
 #include "openvino/pass/graph_rewrite.hpp"
@@ -244,4 +247,63 @@ bool DetectAttentionMask::run_on_model(const std::shared_ptr<ov::Model>& model)
     return false;
 }
 
+// Traces the mask input of Add(QK, mask) backward using BFS.
+// Returns SlidingWindow if a Greater node appears directly inside a
+// BitwiseAnd / BitwiseOr / LogicalAnd (window-size check), Causal otherwise.
+static MaskInfo::MaskType detect_sdpa_mask_type_from_add(const std::shared_ptr<ov::Node>& add_node) {
+    if (!add_node || add_node->get_input_size() < 2) {
+        return MaskInfo::MaskType::Unknown;
+    }
+
+    // Trace mask input (input 1 of Add) backward using BFS.
+    // Look for SWA indicators: a BitwiseAnd / BitwiseOr / LogicalAnd that has
+    // a Greater node as a direct input (window-size check).
+    std::unordered_set<ov::Node*> visited;
+    std::queue<std::shared_ptr<ov::Node>> queue;
+    queue.push(add_node->get_input_node_shared_ptr(1));
+
+    while (!queue.empty()) {
+        auto node = queue.front();
+        queue.pop();
+        if (!node || !visited.insert(node.get()).second)
+            continue;
+
+        // SWA anchor: BitwiseAnd / BitwiseOr / LogicalAnd whose direct input is Greater
+        if (ov::is_type<ov::op::v13::BitwiseAnd>(node) || ov::is_type<ov::op::v13::BitwiseOr>(node) ||
+            ov::is_type<ov::op::v1::LogicalAnd>(node)) {
+            for (size_t i = 0; i < node->get_input_size(); ++i) {
+                if (ov::is_type<ov::op::v1::Greater>(node->get_input_node_shared_ptr(i))) {
+                    return MaskInfo::MaskType::SlidingWindow;
+                }
+            }
+        }
+
+        // Don't cross Parameters or Constants – they are leaf nodes.
+        if (ov::op::util::is_parameter(node) || ov::op::util::is_constant(node))
+            continue;
+
+        for (size_t i = 0; i < node->get_input_size(); ++i)
+            queue.push(node->get_input_node_shared_ptr(i));
+    }
+
+    // No SWA pattern found, treat as causal.
+    return MaskInfo::MaskType::Causal;
+}
+
+bool AnnotatePerSDPAMaskType::run_on_model(const std::shared_ptr<ov::Model>& model) {
+    m_annotations.clear();
+
+    const auto all_patterns = ov::npuw::util::find_all_sdpa_pattern_nodes(model);
+    m_annotations.reserve(all_patterns.size());
+
+    for (const auto& pattern : all_patterns) {
+        if (!pattern.add_node)
+            continue;
+        const auto mask_type = detect_sdpa_mask_type_from_add(pattern.add_node);
+        pattern.add_node->get_rt_info()[NPUW_SDPA_MASK_TYPE_RT_KEY] = static_cast<int>(mask_type);
+        m_annotations.push_back({pattern.add_node->get_friendly_name(), mask_type});
+    }
+    return false;
+}
+
 }  // namespace ov::npuw

src/plugins/intel_npu/src/plugin/npuw/npuw_transformations/detect_causal_mask.hpp

@@ -4,6 +4,9 @@
 
 #pragma once
 
+#include <string>
+#include <vector>
+
 #include "openvino/pass/pass.hpp"
 
 namespace ov::npuw {
@@ -39,4 +42,46 @@ class DetectAttentionMask : public ov::pass::ModelPass {
     MaskInfo m_mask_info;
 };
 
+// rt_info key written by AnnotatePerSDPAMaskType and read by HostFlashAttention::from().
+// Value type: int, corresponding to MaskInfo::MaskType.
+static constexpr const char* NPUW_SDPA_MASK_TYPE_RT_KEY = "npuw_sdpa_mask_type";
+
+// Pre-partitioning pass: annotates each decomposed-SDPA's Add(QK, mask) node in the
+// model with its individual mask type via rt_info[NPUW_SDPA_MASK_TYPE_RT_KEY].
+//
+// This enables per-layer mask-skipping decisions inside HostFlashAttention::from()
+// for mixed SWA + global-attention models (e.g. Gemma-4 E2B/E4B): global-attention
+// ATTN subgraphs can keep mask skipping enabled even when SWA layers are present.
+//
+// Must be run on the whole model BEFORE partitioning so the annotation is carried
+// into the isolated ATTN subgraphs (the Add node object is shared, not cloned).
+// Never modifies the graph structure; run_on_model always returns false.
+class AnnotatePerSDPAMaskType : public ov::pass::ModelPass {
+public:
+    struct Annotation {
+        std::string add_node_name;
+        MaskInfo::MaskType mask_type = MaskInfo::MaskType::Unknown;
+    };
+
+    OPENVINO_MODEL_PASS_RTTI("ov::npuw::AnnotatePerSDPAMaskType");
+    bool run_on_model(const std::shared_ptr<ov::Model>& model) override;
+
+    // Collected per-SDPA mask types from the most recent run_on_model() call.
+    const std::vector<Annotation>& get_annotations() const {
+        return m_annotations;
+    }
+
+    // Convenience helper: returns only mask types in traversal order.
+    std::vector<MaskInfo::MaskType> get_mask_types() const {
+        std::vector<MaskInfo::MaskType> mask_types;
+        mask_types.reserve(m_annotations.size());
+        for (const auto& annotation : m_annotations)
+            mask_types.push_back(annotation.mask_type);
+        return mask_types;
+    }
+
+private:
+    std::vector<Annotation> m_annotations;
+};
+
 }  // namespace ov::npuw

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

@@ -9,6 +9,7 @@
 #include <memory>
 #include <string>
 
+#include "npuw_transformations/detect_causal_mask.hpp"
 #include "openvino/op/add.hpp"
 #include "openvino/op/concat.hpp"
 #include "openvino/op/convert.hpp"
@@ -310,6 +311,53 @@ TEST(HostFlashAttentionFromTest, Fused_MaskTileAtIndexSixInRegularTileWhenMaskSk
     expect_input_name(result->_final_tile_model, 6, "MASK_TILE", "fused final tile");
 }
 
+TEST(HostFlashAttentionFromTest, Fused_PerSDPACausalRtInfo_OverridesGlobalNoAndSkipsRegularMask) {
+    auto model = build_sdpa_model();
+    ASSERT_NE(model, nullptr);
+
+    std::shared_ptr<ov::op::v1::Add> add;
+    for (const auto& node : model->get_ops()) {
+        add = ov::as_type_ptr<ov::op::v1::Add>(node);
+        if (add && add->get_friendly_name() == "add.0")
+            break;
+    }
+    ASSERT_NE(add, nullptr);
+
+    add->get_rt_info()[ov::npuw::NPUW_SDPA_MASK_TYPE_RT_KEY] = static_cast<int>(ov::npuw::MaskInfo::MaskType::Causal);
+
+    // Emulate mixed-model global decision: global NO, but this ATTN subgraph is global/causal.
+    auto result = ov::npuw::function::HostFlashAttention::from(model, true, false);
+    ASSERT_TRUE(result.has_value());
+
+    // Regular tile skips mask (6 inputs), final tile still keeps mask (7 inputs).
+    EXPECT_EQ(result->_tile_model->inputs().size(), 6u);
+    EXPECT_EQ(result->_final_tile_model->inputs().size(), 7u);
+}
+
+TEST(HostFlashAttentionFromTest, Fused_PerSDPASlidingRtInfo_KeepsMaskWhenGlobalNo) {
+    auto model = build_sdpa_model();
+    ASSERT_NE(model, nullptr);
+
+    std::shared_ptr<ov::op::v1::Add> add;
+    for (const auto& node : model->get_ops()) {
+        add = ov::as_type_ptr<ov::op::v1::Add>(node);
+        if (add && add->get_friendly_name() == "add.0")
+            break;
+    }
+    ASSERT_NE(add, nullptr);
+
+    add->get_rt_info()[ov::npuw::NPUW_SDPA_MASK_TYPE_RT_KEY] =
+        static_cast<int>(ov::npuw::MaskInfo::MaskType::SlidingWindow);
+
+    // Emulate mixed-model global decision: global NO and this ATTN subgraph is SWA.
+    auto result = ov::npuw::function::HostFlashAttention::from(model, true, false);
+    ASSERT_TRUE(result.has_value());
+
+    // Regular tile keeps mask (7 inputs), final tile always keeps mask (7 inputs).
+    EXPECT_EQ(result->_tile_model->inputs().size(), 7u);
+    EXPECT_EQ(result->_final_tile_model->inputs().size(), 7u);
+}
+
 // ============================================================================
 // Tile param index map
 // ============================================================================