diff --git a/src/plugins/intel_cpu/src/nodes/executors/x64/subgraph.hpp b/src/plugins/intel_cpu/src/nodes/executors/x64/subgraph.hpp index c8e9b0dd2f55a1..dd72ce552db43e 100644 --- a/src/plugins/intel_cpu/src/nodes/executors/x64/subgraph.hpp +++ b/src/plugins/intel_cpu/src/nodes/executors/x64/subgraph.hpp @@ -29,13 +29,7 @@ class SubgraphExecutor : public SubgraphRepackingExecutor { const BufferScratchpadAllocator& allocator, const ov::intel_cpu::MultiCacheWeakPtr& kernel_cache); -#ifdef SNIPPETS_DEBUG_CAPS -protected: - void segfault_detector() const override; -private: - bool enabled_segfault_detector = false; -#endif }; class SubgraphStaticExecutor : public SubgraphRepackingStaticExecutor { diff --git a/src/plugins/intel_npu/src/al/include/intel_npu/config/npuw_option_defs.inc b/src/plugins/intel_npu/src/al/include/intel_npu/config/npuw_option_defs.inc index 530cd6d5f7ccd2..f1cf27613fe282 100644 --- a/src/plugins/intel_npu/src/al/include/intel_npu/config/npuw_option_defs.inc +++ b/src/plugins/intel_npu/src/al/include/intel_npu/config/npuw_option_defs.inc @@ -73,6 +73,8 @@ INTEL_NPU_NPUW_SIMPLE_OPT(NPUW_LLM_PREFIX_CACHING_MAX_NUM_BLOCKS, uint64_t, 128, INTEL_NPU_NPUW_SIMPLE_OPT(NPUW_LLM_ENABLE_CONTINUOUS_PREFILL, bool, false, ov::intel_npu::npuw::llm, enable_continuous_prefill, "NPUW_LLM_ENABLE_CONTINUOUS_PREFILL", LLM, EXPOSED, CACHED, ALL) INTEL_NPU_NPUW_SIMPLE_OPT(NPUW_LLM_CACHE_ROPE, bool, true, ov::intel_npu::npuw::llm, cache_rope, "NPUW_LLM_CACHE_ROPE", LLM, EXPOSED, CACHED, ALL) INTEL_NPU_NPUW_SIMPLE_OPT(NPUW_LLM_ENABLE_BLOCK_BASED_KV_CACHE, bool, false, ov::intel_npu::npuw::llm, enable_block_based_kv_cache, "NPUW_LLM_ENABLE_BLOCK_BASED_KV_CACHE", LLM, EXPOSED, CACHED, ALL) +INTEL_NPU_NPUW_SIMPLE_OPT(NPUW_LLM_SLIDING_WINDOW, uint64_t, 0, ov::intel_npu::npuw::llm, sliding_window, "NPUW_LLM_SLIDING_WINDOW", LLM, EXPOSED, CACHED, ALL) +INTEL_NPU_NPUW_SIMPLE_OPT(NPUW_LLM_LAYER_TYPES, std::string, "", ov::intel_npu::npuw::llm, layer_types, "NPUW_LLM_LAYER_TYPES", LLM, EXPOSED, CACHED, ALL) INTEL_NPU_NPUW_STRING_ENUM_OPT(NPUW_LLM_PREFILL_MOE_HINT, ::intel_npu::npuw::llm::MoEHint, MoEHintOptionTraits, ov::intel_npu::npuw::llm, prefill_moe_hint, "NPUW_LLM_PREFILL_MOE_HINT", LLM, EXPOSED, CACHED, ALL) INTEL_NPU_NPUW_STRING_ENUM_OPT(NPUW_LLM_GENERATE_MOE_HINT, ::intel_npu::npuw::llm::MoEHint, MoEHintOptionTraits, ov::intel_npu::npuw::llm, generate_moe_hint, "NPUW_LLM_GENERATE_MOE_HINT", LLM, EXPOSED, CACHED, ALL) INTEL_NPU_NPUW_STRING_ENUM_OPT(NPUW_LLM_PREFILL_HINT, ::intel_npu::npuw::llm::PrefillHint, PrefillHintOptionTraits, ov::intel_npu::npuw::llm, prefill_hint, "NPUW_LLM_PREFILL_HINT", LLM, EXPOSED, CACHED, ALL) 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 4d42da3fe027ae..4d81a13e3caa3b 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 @@ -882,30 +882,43 @@ ov::npuw::v1::subgraphs::RuntimeBehaviorFactory make_runtime_factory() { OPENVINO_ASSERT(hfa_desc->is_valid(), "HFA configuration must be valid"); const int64_t tile_size = hfa_desc->_tile_size; - const int64_t total_kv_length = state.hfa_selector->context_length(); - const int64_t num_tiles = total_kv_length / tile_size; - OPENVINO_ASSERT(total_kv_length % tile_size == 0, - "HFA total KV length must be multiple of tile size for now"); const auto& hfa_inputs = io.inputs; const auto& sdpa_info = hfa_desc->_sdpa_attention_info; const auto& sdpa_in = sdpa_info._sdpa_indices; + const uint32_t K_SEQ_DIM = static_cast(sdpa_info._k_seq_dim); + const uint32_t V_SEQ_DIM = static_cast(sdpa_info._v_seq_dim); - // Collect all KV block tensors (works for single-block and multi-block cases) + // Collect all KV block tensors and accumulate block_sum in one pass. + // (works for single-block and multi-block cases) NPUW_ASSERT(!sdpa_in.past_key_blocks.empty() && !sdpa_in.past_value_blocks.empty() && "SDPA indices must have at least one past_key/value block"); NPUW_ASSERT(sdpa_in.past_key_blocks.size() == sdpa_in.past_value_blocks.size() && "Number of past key blocks must match number of past value blocks"); std::vector> past_key_blocks; std::vector> past_value_blocks; + int64_t block_sum = 0; for (size_t i = 0; i < sdpa_in.past_key_blocks.size(); ++i) { past_key_blocks.push_back(hfa_inputs.at(sdpa_in.past_key_blocks[i])); past_value_blocks.push_back(hfa_inputs.at(sdpa_in.past_value_blocks[i])); + block_sum += static_cast(past_key_blocks.back()->get_shape()[K_SEQ_DIM]); } auto query_tensor = hfa_inputs.at(sdpa_in.query); auto present_key_tensor = hfa_inputs.at(sdpa_in.present_key); auto attention_mask_tensor = hfa_inputs.at(sdpa_in.attention_mask); auto present_value_tensor = hfa_inputs.at(sdpa_in.present_value); + block_sum += static_cast(present_key_tensor->get_shape()[K_SEQ_DIM]); + + // total_kv_length = min(context_length(), block_sum) works for both: + // Global SDPA: block_sum == context_length() → min picks either. + // SWA within window: block_sum == context_length() → same. + // SWA past window: block_sum < context_length() (blocks capped at window_size) + // → min picks block_sum, the actual occupied KV length. + const int64_t total_kv_length = std::min(state.hfa_selector->context_length(), block_sum); + const int64_t num_tiles = total_kv_length / tile_size; + OPENVINO_ASSERT(total_kv_length % tile_size == 0, + "HFA total KV length must be multiple of tile size for now"); + auto& regular_tile_request = state.hfa_requests.infer_requests[HFARequestSet::REGULAR_TILE]; auto& final_tile_request = state.hfa_requests.infer_requests[HFARequestSet::FINAL_TILE]; auto attention_output_tensor = @@ -956,8 +969,6 @@ ov::npuw::v1::subgraphs::RuntimeBehaviorFactory make_runtime_factory() { final_tile_request->set_tensor(hfa_desc->_compiled_final_tile_model->outputs()[0], attention_output_tensor); - const uint32_t K_SEQ_DIM = static_cast(sdpa_info._k_seq_dim); - const uint32_t V_SEQ_DIM = static_cast(sdpa_info._v_seq_dim); constexpr uint32_t MASK_KV_SEQ_DIM = 3; size_t next_available_mask_buffer_idx = 0; diff --git a/src/plugins/intel_npu/src/plugin/npuw/host_flash_attention.cpp b/src/plugins/intel_npu/src/plugin/npuw/host_flash_attention.cpp index 87228370eaa18d..2d0ca99644d0b5 100644 --- a/src/plugins/intel_npu/src/plugin/npuw/host_flash_attention.cpp +++ b/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" @@ -659,6 +660,10 @@ static std::shared_ptr create_hfa_tile_model(const ov::Shape& q_shape // enable_mask_skipping is true (depending on the model mask type). // For the non-fused operation all tiles require mask const bool use_mask = is_final_tile || !fused_flash_attention || !enable_mask_skipping; + LOG_INFO("[HFA] " << (is_final_tile ? "FINAL" : "regular") << " tile mask decision: use_mask=" << use_mask + << " (is_final_tile=" << is_final_tile + << ", fused_flash_attention=" << fused_flash_attention + << ", enable_mask_skipping=" << enable_mask_skipping << ")"); auto f32_nodes = convert_inputs_to_f32(inputs, mask_dtype, compute_dtype, use_mask); FlashAttentionResults results; @@ -977,6 +982,9 @@ std::optional HostFlashAttention::from(const std::shared_ptr bool enable_mask_skipping) { LOG_INFO("Attempting to create HostFlashAttention" << (fused_flash_attention ? " with fused flash attention node" : "")); + LOG_INFO("[HFA] enable_mask_skipping=" << enable_mask_skipping + << " (regular, non-final tiles will " + << (enable_mask_skipping ? "SKIP" : "KEEP") << " the explicit mask)"); LOG_BLOCK(); // ======================================================================== @@ -1106,6 +1114,32 @@ std::optional 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(it->second.as()); + 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 +1149,7 @@ std::optional 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 +1165,7 @@ std::optional 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) { diff --git a/src/plugins/intel_npu/src/plugin/npuw/infer_request_utils.cpp b/src/plugins/intel_npu/src/plugin/npuw/infer_request_utils.cpp index bd2a2c91eabd7e..75fff08aa53e64 100644 --- a/src/plugins/intel_npu/src/plugin/npuw/infer_request_utils.cpp +++ b/src/plugins/intel_npu/src/plugin/npuw/infer_request_utils.cpp @@ -4,6 +4,7 @@ #include "infer_request_utils.hpp" +#include #include #include "logging.hpp" @@ -150,6 +151,125 @@ void ov::npuw::util::copy_tensor_by_dim(ov::SoPtr src_tensor, } } +void ov::npuw::util::write_kv_slice_sliding(ov::SoPtr dst_tensor, + ov::SoPtr src_new_kv, + uint32_t dst_kv_dim, + uint32_t src_kv_dim, + uint32_t num_stored_tokens_before, + uint32_t num_new_tokens, + SlidingBufferLayout layout) { + const uint32_t capacity = static_cast(dst_tensor->get_shape()[dst_kv_dim]); + const uint32_t old_total = num_stored_tokens_before; + const uint32_t new_total = old_total + num_new_tokens; + const uint32_t old_valid = std::min(old_total, capacity); + const uint32_t new_valid = std::min(new_total, capacity); + + // Clamp against the source's own length too: a source tensor may legitimately hold + // fewer valid tokens than `num_new_tokens` claims (e.g. when re-using another + // layer's already-capacity-limited past buffer as a source, see the header comment). + const uint32_t src_len = static_cast(src_new_kv->get_shape()[src_kv_dim]); + const uint32_t tokens_to_write = std::min({num_new_tokens, new_valid, src_len}); + + if (layout == SlidingBufferLayout::Circular) { + // No shift, ever: token at absolute position p always lives at physical index + // (p % capacity). See SlidingBufferLayout's doc comment in the header for why + // this is safe. Skip the leading tokens of this call that would be immediately + // overwritten later in the very same call (mirrors the LeftAligned clamp above). + if (tokens_to_write == 0) { + return; + } + const uint32_t first_new_abs_pos = num_stored_tokens_before + (num_new_tokens - tokens_to_write); + const uint32_t dst_start = first_new_abs_pos % capacity; + + auto src_slice = (src_len > tokens_to_write) + ? make_tensor_slice(src_new_kv, src_kv_dim, src_len - tokens_to_write, src_len) + : src_new_kv; + + if (dst_start + tokens_to_write <= capacity) { + // Single contiguous write - also covers the not-yet-saturated warm-up + // phase, where dst_start == first_new_abs_pos, i.e. a plain append. + auto dst_slice = make_tensor_slice(dst_tensor, dst_kv_dim, dst_start, dst_start + tokens_to_write); + copy_tensor_by_dim(src_slice, dst_slice, src_kv_dim, dst_kv_dim); + } else { + // Wraps past the end of the buffer: split into two contiguous legs. + const uint32_t first_leg_len = capacity - dst_start; + const uint32_t second_leg_len = tokens_to_write - first_leg_len; + + auto src_first_leg = make_tensor_slice(src_slice, src_kv_dim, 0u, first_leg_len); + auto dst_first_leg = make_tensor_slice(dst_tensor, dst_kv_dim, dst_start, capacity); + copy_tensor_by_dim(src_first_leg, dst_first_leg, src_kv_dim, dst_kv_dim); + + auto src_second_leg = make_tensor_slice(src_slice, src_kv_dim, first_leg_len, tokens_to_write); + auto dst_second_leg = make_tensor_slice(dst_tensor, dst_kv_dim, 0u, second_leg_len); + copy_tensor_by_dim(src_second_leg, dst_second_leg, src_kv_dim, dst_kv_dim); + } + return; + } + + const uint32_t keep = new_valid - tokens_to_write; + const bool needs_shift = (keep > 0 && keep < old_valid); + + if (needs_shift && dst_kv_dim == 3u) { + // Transposed-V layout (dst_kv_dim == 3): a partial-slice shift touches only + // `old_valid` of the `capacity` columns, but a dim-3 slice of a [1,C,H,W] tensor + // is non-contiguous, so both the read (old_tail->copy_to) and the write + // (copy_tensor_by_dim -> copy_columns_by_row_chunks) legs degrade into C*H + // individual small (per-token) memory transactions. When dst_tensor lives in + // NPU-resident remote memory, per-transaction latency (not bytes moved) + // dominates, and C*H can be in the thousands - this is the empirically + // dominant cost of the sliding-window KV update (~600ms/step on real HW). + // + // Since the *whole* (unsliced) buffer is fully contiguous, round-trip it as a + // single big contiguous transfer instead: one bulk device->CPU copy, a cheap + // in-CPU-memory shift (regular DRAM, C*H iterations here are negligible), then + // one bulk CPU->device copy back. This trades "only move what changed" for + // "always move `capacity` columns, but in O(1) device-memory transactions". + LOG_DEBUG("[SWA] Bulk-shifting KV buffer (dim=3): keeping last " + << keep << " of " << old_valid << " old token(s), capacity=" << capacity); + auto whole_tmp = allocMem(dst_tensor->get_element_type(), dst_tensor->get_shape(), "CPU", nullptr); + dst_tensor->copy_to(whole_tmp._ptr); // single bulk contiguous transfer + + auto old_tail_cpu = make_tensor_slice(whole_tmp, dst_kv_dim, old_valid - keep, old_valid); + auto shift_tmp = allocMem(dst_tensor->get_element_type(), old_tail_cpu->get_shape(), "CPU", nullptr); + old_tail_cpu->copy_to(shift_tmp._ptr); // CPU-to-CPU, cheap regardless of iteration count + auto dst_front_cpu = make_tensor_slice(whole_tmp, dst_kv_dim, 0u, keep); + copy_tensor_by_dim(shift_tmp, dst_front_cpu, dst_kv_dim, dst_kv_dim); // CPU-to-CPU + + if (tokens_to_write > 0) { + auto src_slice = (src_len > tokens_to_write) + ? make_tensor_slice(src_new_kv, src_kv_dim, src_len - tokens_to_write, src_len) + : src_new_kv; + auto dst_back_cpu = make_tensor_slice(whole_tmp, dst_kv_dim, keep, keep + tokens_to_write); + copy_tensor_by_dim(src_slice, dst_back_cpu, src_kv_dim, dst_kv_dim); + } + + whole_tmp->copy_to(dst_tensor._ptr); // single bulk contiguous transfer back + return; + } + + if (needs_shift) { + // Sliding window is (re)saturated: shift the surviving tail of the old content to + // the front of the buffer. A temporary CPU snapshot is used because dst and the + // "old" region alias the same tensor, making a direct in-place copy unsafe. + LOG_DEBUG("[SWA] Shifting KV buffer: keeping last " << keep << " of " << old_valid << " old token(s), dim=" + << dst_kv_dim << ", capacity=" << capacity); + auto old_tail = make_tensor_slice(dst_tensor, dst_kv_dim, old_valid - keep, old_valid); + auto tmp = allocMem(dst_tensor->get_element_type(), old_tail->get_shape(), "CPU", nullptr); + old_tail->copy_to(tmp._ptr); + auto dst_front = make_tensor_slice(dst_tensor, dst_kv_dim, 0u, keep); + copy_tensor_by_dim(tmp, dst_front, dst_kv_dim, dst_kv_dim); + } + + if (tokens_to_write == 0) { + return; + } + auto src_slice = (src_len > tokens_to_write) + ? make_tensor_slice(src_new_kv, src_kv_dim, src_len - tokens_to_write, src_len) + : src_new_kv; + auto dst_back = make_tensor_slice(dst_tensor, dst_kv_dim, keep, keep + tokens_to_write); + copy_tensor_by_dim(src_slice, dst_back, src_kv_dim, dst_kv_dim); +} + std::optional> ov::npuw::util::find_port_by_name( const std::vector>& ports, const std::string& name) { diff --git a/src/plugins/intel_npu/src/plugin/npuw/infer_request_utils.hpp b/src/plugins/intel_npu/src/plugin/npuw/infer_request_utils.hpp index e79c34a352c864..1da4f926608788 100644 --- a/src/plugins/intel_npu/src/plugin/npuw/infer_request_utils.hpp +++ b/src/plugins/intel_npu/src/plugin/npuw/infer_request_utils.hpp @@ -33,6 +33,55 @@ void copy_tensor_by_dim(ov::SoPtr src_tensor, uint32_t kv_dim_src, uint32_t kv_dim_dst); +// Physical layout maintained by write_kv_slice_sliding() for a saturated (capacity-C, +// total_tokens > C) sliding-window buffer: +// - LeftAligned: the most recent C tokens are kept left-aligned at [0, C), in strict +// chronological order. Reaching this state requires shifting the surviving tail +// forward every time the window (re)saturates. +// - Circular: token at absolute position p always lives at physical index (p % C) - +// no data is ever moved, only overwritten in place. This is mathematically +// equivalent for callers whose only consumers are (a) the compiled model's mask, +// which gates the past-KV region with an unconditional "structural index < +// past_kv_len" check (no per-column temporal identity check - see +// rebuild_sliding_window_mask() in sliding_window_mask.cpp) and (b) attention +// itself, which is a permutation-invariant reduction over the visible K/V columns. +// RoPE is baked into K once at write time (using the token's own true absolute +// position), never recomputed from buffer physical position at read time, so +// reordering physical slots doesn't affect correctness either. +// Circular MUST NOT be used for a buffer that is ever read back elsewhere as a +// plain contiguous/left-aligned source (e.g. Continuous strategy's variant-switch +// migration or chunked-prefill past-KV reuse) unless that reader is updated to +// unwrap the circular layout first. +enum class SlidingBufferLayout { LeftAligned, Circular }; + +// Writes `src_new_kv` (holding the freshly-produced KV content, of which the last +// `num_new_tokens` entries along `src_kv_dim` are meaningful) into `dst_tensor`'s +// past-KV buffer along `dst_kv_dim`, honoring `dst_tensor`'s own capacity +// (dst_tensor->get_shape()[dst_kv_dim]) - which may be smaller than the logical +// "total tokens seen so far" for Sliding Window Attention (SWA) layers, since their +// past_key_values Parameter is reshaped to the (smaller) window size at compile time. +// +// `layout` selects the physical arrangement used once the window (re)saturates - see +// SlidingBufferLayout above. Defaults to LeftAligned, matching all pre-existing +// callers' expectations. +// +// For non-SWA layers (capacity >= total tokens ever seen), this is equivalent to the +// original unconditional "write at [old_total, new_total)" behavior in both layouts - +// no shifting/wrapping ever occurs and this call is a drop-in replacement. +// +// NB: `src_new_kv` is expected to follow the "present/output" convention - its +// meaningful content is right-aligned at the tail (mirrors update_kvcache_for's +// existing src_seq_len > num_tokens handling). Persistent *past* buffers reused as a +// source (e.g. chunked-prefill's own past_key_values) are LEFT-aligned instead and +// must be pre-sliced by the caller to their valid prefix before being passed in here. +void write_kv_slice_sliding(ov::SoPtr dst_tensor, + ov::SoPtr src_new_kv, + uint32_t dst_kv_dim, + uint32_t src_kv_dim, + uint32_t num_stored_tokens_before, + uint32_t num_new_tokens, + SlidingBufferLayout layout = SlidingBufferLayout::LeftAligned); + std::optional> find_port_by_name(const std::vector>& ports, const std::string& name); /** diff --git a/src/plugins/intel_npu/src/plugin/npuw/kv_cache_block_manager.cpp b/src/plugins/intel_npu/src/plugin/npuw/kv_cache_block_manager.cpp index df01095ff4958a..9d617772c82201 100644 --- a/src/plugins/intel_npu/src/plugin/npuw/kv_cache_block_manager.cpp +++ b/src/plugins/intel_npu/src/plugin/npuw/kv_cache_block_manager.cpp @@ -77,6 +77,7 @@ std::optional KVCacheBlockManager::allocate_block() { // Reset block state block.num_tokens = 0; block.is_allocated = true; + allocated_order_.push_back(block_id); LOG_VERB("KVCacheBlockManager: Allocated block " << block_id << " (free blocks remaining: " << free_block_ids_.size() << ")"); @@ -130,29 +131,62 @@ uint32_t KVCacheBlockManager::get_block_tokens(uint32_t block_id) const { } std::vector KVCacheBlockManager::get_allocated_blocks() const { - std::vector allocated; - allocated.reserve(max_blocks_ - free_block_ids_.size()); + return std::vector(allocated_order_.begin(), allocated_order_.end()); +} - for (uint32_t i = 0; i < blocks_.size(); ++i) { - if (blocks_[i].is_allocated) { - allocated.push_back(i); +void KVCacheBlockManager::ensure_blocks_up_to(uint32_t up_to_index) { + // evicted_count_ + allocated_order_.size() is the total number of chronological + // blocks ever allocated by this manager - keep allocating (evicting the oldest + // resident block once the pool is full) until it covers up_to_index. + // + // NB: when the pool is already full, this is the EXPECTED steady-state path for a + // sliding-window layer (happens once per new chronological block, i.e. once every + // block_size tokens) - so eviction is done directly here, without ever going through + // allocate_block()'s "pool exhausted" warning (that warning is reserved for the + // genuinely-unexpected exhaustion of a non-evicting, e.g. full-attention, manager). + while (evicted_count_ + allocated_order_.size() <= up_to_index) { + if (free_block_ids_.empty()) { + OPENVINO_ASSERT(!allocated_order_.empty(), + "KVCacheBlockManager: pool exhausted (max_blocks=", + max_blocks_, + ") and no resident block available to evict"); + + const uint32_t oldest_id = allocated_order_.front(); + allocated_order_.pop_front(); + auto& oldest_block = blocks_[oldest_id]; + oldest_block.is_allocated = false; + oldest_block.num_tokens = 0; + free_block_ids_.push(oldest_id); + ++evicted_count_; + LOG_VERB("KVCacheBlockManager: Evicted oldest block " << oldest_id << " (evicted_count=" << evicted_count_ + << ")"); } + + const auto block_id = allocate_block(); + OPENVINO_ASSERT(block_id.has_value(), + "KVCacheBlockManager: unexpected allocate_block() failure right after eviction"); } +} - return allocated; +uint32_t KVCacheBlockManager::resident_index(uint32_t chronological_index) const { + OPENVINO_ASSERT(chronological_index >= evicted_count_, + "KVCacheBlockManager: block ", + chronological_index, + " has already been evicted (evicted_count=", + evicted_count_, + ")"); + return chronological_index - evicted_count_; } void KVCacheBlockManager::release(uint32_t keep_warm_count) { LOG_DEBUG("KVCacheBlockManager: Resetting blocks (keep_warm_count=" << keep_warm_count << ")"); - // Iterate allocated blocks in ascending block-ID order. + // Iterate allocated blocks in allocation order (oldest first). // The first keep_warm_count retain their device tensors (warm reuse next round). // The remainder have tensors dropped to reduce RSS. uint32_t warm_seen = 0; - for (auto& block : blocks_) { - if (!block.is_allocated) { - continue; - } + for (uint32_t block_id : allocated_order_) { + auto& block = blocks_[block_id]; block.is_allocated = false; block.num_tokens = 0; if (warm_seen >= keep_warm_count) { @@ -166,6 +200,8 @@ void KVCacheBlockManager::release(uint32_t keep_warm_count) { << keep_warm_count << " warm"); } + allocated_order_.clear(); + evicted_count_ = 0; rebuild_free_block_ids(); } @@ -187,24 +223,25 @@ void KVCacheBlockManager::rebuild_free_block_ids() { } void KVCacheBlockManager::truncate_allocated(uint32_t keep_count) { - const auto allocated = get_allocated_blocks(); - OPENVINO_ASSERT(keep_count <= allocated.size(), + OPENVINO_ASSERT(keep_count <= allocated_order_.size(), "KVCacheBlockManager: cannot truncate to ", keep_count, " blocks, only ", - allocated.size(), + allocated_order_.size(), " are allocated"); // Deallocate the suffix, keeping device tensors warm for quick reuse. - for (size_t i = keep_count; i < allocated.size(); ++i) { - auto& block = blocks_[allocated[i]]; + const auto dropped_count = allocated_order_.size() - keep_count; + for (size_t i = keep_count; i < allocated_order_.size(); ++i) { + auto& block = blocks_[allocated_order_[i]]; block.is_allocated = false; block.num_tokens = 0; } + allocated_order_.resize(keep_count); rebuild_free_block_ids(); - LOG_DEBUG("KVCacheBlockManager: truncated to " << keep_count << " blocks, " << (allocated.size() - keep_count) + LOG_DEBUG("KVCacheBlockManager: truncated to " << keep_count << " blocks, " << dropped_count << " suffix block(s) freed"); } diff --git a/src/plugins/intel_npu/src/plugin/npuw/kv_cache_block_manager.hpp b/src/plugins/intel_npu/src/plugin/npuw/kv_cache_block_manager.hpp index c07768ff6b315e..d015c4e690749b 100644 --- a/src/plugins/intel_npu/src/plugin/npuw/kv_cache_block_manager.hpp +++ b/src/plugins/intel_npu/src/plugin/npuw/kv_cache_block_manager.hpp @@ -5,6 +5,7 @@ #pragma once #include +#include #include #include #include @@ -97,12 +98,52 @@ class KVCacheBlockManager { uint32_t get_block_tokens(uint32_t block_id) const; /** - * @brief Get list of all currently allocated block IDs + * @brief Get list of all currently allocated block IDs, in allocation order + * (oldest first). This is the order new blocks were appended in, which - unlike + * ascending block ID order - stays correct even after ensure_blocks_up_to() has + * reused evicted IDs for newer blocks. * * @return Vector of block IDs */ std::vector get_allocated_blocks() const; + /** + * @brief Ensure blocks exist for every chronological (ever-increasing) block index + * up to and including `up_to_index`, allocating new blocks as needed. + * + * Intended for sliding-window (SWA) layers: once max_blocks are resident, each + * additional index evicts the single oldest resident block (by allocation order, + * see get_allocated_blocks()) and reuses its slot, so the manager always holds + * exactly the most recent max_blocks chronological blocks. Use resident_index() / + * evicted_count() to translate a chronological index into its position in + * get_allocated_blocks() / get_allocated_blocks().size(). + * + * Callers that never need eviction (full-attention layers) should keep using + * allocate_block() directly, which fails loudly on exhaustion instead of evicting. + * + * @param up_to_index Highest chronological block index that must be resident. + * Indices must be requested in non-decreasing order across the + * manager's lifetime (matches token generation order). + */ + void ensure_blocks_up_to(uint32_t up_to_index); + + /** + * @brief Number of chronological blocks evicted so far via ensure_blocks_up_to(). + * Always 0 unless ensure_blocks_up_to() has actually evicted a block. + */ + uint32_t evicted_count() const { + return evicted_count_; + } + + /** + * @brief Translate a chronological block index into its position in + * get_allocated_blocks() (0 = oldest resident block). + * + * @param chronological_index Must be >= evicted_count(), i.e. still resident - + * indices below that have been evicted and no longer have a backing block. + */ + uint32_t resident_index(uint32_t chronological_index) const; + /** * @brief Release blocks to FREE state, selectively dropping device memory. * @@ -177,6 +218,15 @@ class KVCacheBlockManager { std::vector blocks_; ///< All blocks (free + allocated) std::stack free_block_ids_; ///< Stack of free block IDs (LIFO for better reuse) + /// Allocated block IDs in allocation order (oldest first). Backs get_allocated_blocks() + /// and drives ensure_blocks_up_to()'s eviction (front = oldest = evicted first). + /// For managers that never evict, this always equals ascending block-ID order. + std::deque allocated_order_; + + /// Number of chronological blocks evicted so far via ensure_blocks_up_to() (0 unless + /// eviction has actually happened). + uint32_t evicted_count_ = 0; + ov::element::Type element_type_; ///< Element type for tensors ov::Shape block_shape_; ///< Shape for block tensors std::string device_; ///< Target device diff --git a/src/plugins/intel_npu/src/plugin/npuw/llm_block_kvcache_strategy.cpp b/src/plugins/intel_npu/src/plugin/npuw/llm_block_kvcache_strategy.cpp index 8859dbce0b87e8..67f4c64cbfb19a 100644 --- a/src/plugins/intel_npu/src/plugin/npuw/llm_block_kvcache_strategy.cpp +++ b/src/plugins/intel_npu/src/plugin/npuw/llm_block_kvcache_strategy.cpp @@ -12,6 +12,7 @@ #include "llm_infer_base_request.hpp" #include "llm_infer_request.hpp" #include "logging.hpp" +#include "perf.hpp" #include "util.hpp" namespace { @@ -252,6 +253,11 @@ void LLMBlockKVCacheStrategy::on_initialize() { // Create block managers and pre-compute per-variant binding helpers create_block_managers_and_helpers(); + // SWA layers: bind their numbered ports once as static views into a single persistent + // window buffer (see LayerBlockManagers doc comment). This never needs to be redone + // except after on_reset() overwrites ports with dummy tensors. + bind_swa_window_views(); + // Snapshot original prefill output tensors (for restore_prefill_output_buffers()) and // build m_output_kv_info (output_name → layer/kv) to avoid per-call regex. for (const auto& [layer_idx, layer_managers] : m_kv_cache_block_managers) { @@ -352,6 +358,36 @@ void LLMBlockKVCacheStrategy::on_reset(uint32_t next_prompt_length) { } } } + + // ── Step 5: restore SWA window buffers ─────────────────────────────────────────── + // Step 2 above overwrote EVERY numbered block port (SWA and non-sliding alike) with a + // dummy tensor. SWA layers' window buffers are persistent across conversations (unlike + // non-sliding blocks, they are never released) - re-establish the static views and + // clear stale content from the previous conversation (harmless either way, since + // write_kv_slice_sliding() only ever reads/writes the logically-valid prefix, but kept + // for defense-in-depth against stale reads). + for (auto& [layer_idx, layer_managers] : m_kv_cache_block_managers) { + if (!layer_managers.is_sliding) { + continue; + } + if (layer_managers.swa_key_window) { + ov::npuw::util::fill_tensor_bytes(layer_managers.swa_key_window, 0u); + } + if (layer_managers.swa_value_window) { + ov::npuw::util::fill_tensor_bytes(layer_managers.swa_value_window, 0u); + } + } + bind_swa_window_views(); + + // New conversation: report the previous one's SWA update cost (if any) and reset the + // running perf counters so the next conversation's average isn't skewed by this one. + if (m_swa_update_calls > 0) { + LOG_INFO("[PERF] Previous conversation: update_swa_windows_generate avg=" + << (m_swa_update_total_ms / m_swa_update_calls) << "ms over " << m_swa_update_calls + << " step(s), total=" << m_swa_update_total_ms << "ms"); + } + m_swa_update_total_ms = 0.0f; + m_swa_update_calls = 0; } void LLMBlockKVCacheStrategy::on_prefill_chunk_begin(uint32_t current_prompts_len) { @@ -397,6 +433,12 @@ void LLMBlockKVCacheStrategy::on_prefill_chunk_done(uint32_t current_prompts_len v_transposed, write_start); } + + // SWA layers are never zero-copy-redirected (redirect_prefill_outputs_to_new_blocks() + // skips them - see its process_kv_blocks lambda), so their "present.N.key/value" outputs + // always sit in the prefill model's own (unredirected) buffers regardless of + // m_zero_copy_last_chunk - this call is therefore always safe/correct to run. + update_swa_windows_from_prefill(current_prompts_len); } void LLMBlockKVCacheStrategy::on_generate_kv_init() { @@ -474,12 +516,207 @@ void LLMBlockKVCacheStrategy::on_generate_step_done(uint32_t input_tokens_len) { const auto& kvcache_desc = m_req.m_npuw_llm_compiled_model->m_kvcache_desc; const uint32_t tokens_after = kvcache_desc.num_stored_tokens; const uint32_t tokens_before = tokens_after - input_tokens_len; - copy_outputs_to_blocks(m_req.m_kvcache_request, - m_req.m_kvcache_out_ports, - input_tokens_len, - kvcache_desc.v_tensors_transposed_gen, - tokens_before); - update_generate_bindings(tokens_before, tokens_after, m_req.m_kvcache_request); + + const float copy_ms = ov::npuw::perf::ms_to_run([&]() { + copy_outputs_to_blocks(m_req.m_kvcache_request, + m_req.m_kvcache_out_ports, + input_tokens_len, + kvcache_desc.v_tensors_transposed_gen, + tokens_before); + }); + + const float bindings_ms = ov::npuw::perf::ms_to_run([&]() { + update_generate_bindings(tokens_before, tokens_after, m_req.m_kvcache_request); + }); + + // SWA layers: slide the window buffer forward using this step's freshly computed + // "present.N.key/value" outputs (copy_outputs_to_blocks()/update_generate_bindings() + // above are both no-ops for SWA layers, since they only act on non-null + // key_manager/value_manager). + const float swa_ms = ov::npuw::perf::ms_to_run([&]() { + update_swa_windows_generate(input_tokens_len); + }); + + m_swa_update_total_ms += swa_ms; + m_swa_update_calls++; + LOG_INFO("[PERF] on_generate_step_done: copy_outputs_to_blocks=" << copy_ms << "ms, update_generate_bindings=" + << bindings_ms << "ms, update_swa_windows_generate=" + << swa_ms << "ms (avg=" << (m_swa_update_total_ms / + m_swa_update_calls) << "ms over " << m_swa_update_calls + << " calls, total=" << m_swa_update_total_ms << "ms)"); + // Periodic INFO-level summary so the cost is visible without switching to DEBUG log + // level for a long generation run. + if (m_swa_update_calls % 32 == 0) { + LOG_INFO("[PERF] update_swa_windows_generate: avg=" << (m_swa_update_total_ms / m_swa_update_calls) + << "ms over " << m_swa_update_calls + << " step(s), total=" << m_swa_update_total_ms << "ms"); + } +} + +// ============================================================================ +// Private: SWA window buffer management +// +// SWA layers do not use KVCacheBlockManager: a per-token-exact sliding window cannot be +// represented by whole-block eviction/rotation, because a per-token-exact window cutoff +// (e.g. window_size not a multiple of block_size) can only be reproduced at token +// granularity, not by evicting/rotating whole blocks - block-granularity rebinding would +// either keep too many or too few tokens whenever the window doesn't land on a block +// boundary ("torn" mid-block positions). +// +// Instead, each SWA layer gets ONE persistent, contiguous window_size-wide buffer per +// key/value (swa_key_window / swa_value_window), allocated once in +// create_block_managers_and_helpers(). Its k = window_size/block_size numbered ports are +// bound ONCE (bind_swa_window_views()) as adjacent VIEWS into this single buffer and are +// NEVER rebound again - only the buffer's CONTENT is updated in place, each prefill chunk +// and each generate step, via the shared write_kv_slice_sliding() utility. +// +// IMPORTANT: ports must stay bound to a FIXED address for their entire lifetime. An +// oversized-buffer + periodic-compaction + "slide the view forward every step" scheme was +// tried and reverted: the NPU backend's zero-copy remote tensor binding only supports a +// stable, block-aligned address per port - rebinding to an arbitrary shifting offset every +// step triggered "Strided remote tensor is not supported for this port!" (falling back to a +// CPU strided-copy path) and produced WRONG generation results on real hardware. Any future +// perf work on this code path must keep every port's binding fixed. +// +// Physical layout: the window buffer's content is maintained with +// SlidingBufferLayout::Circular (see infer_request_utils.hpp) instead of the +// LeftAligned/shift-based layout used by the sibling LLMContinuousKVCacheStrategy. Once +// saturated, token at absolute position p always lives at physical index (p % +// window_size) - new tokens simply overwrite their circular slot, no existing content is +// ever shifted or copied elsewhere. This is safe here specifically because this buffer's +// ONLY consumer is the compiled model itself via the fixed zero-copy port bindings above - +// nothing in this strategy ever reads it back as a plain left-aligned source (unlike +// LLMContinuousKVCacheStrategy's variant-switch migration and chunked-prefill past-KV +// reuse, which is why Circular is not used there). Correctness rests on: the compiled +// mask gates the whole past-KV region with an unconditional "structural index < +// past_kv_len" check rather than verifying each column's temporal identity (see +// rebuild_sliding_window_mask() in sliding_window_mask.cpp), and attention is a +// permutation-invariant reduction over the visible K/V columns, so physical column order +// within that region does not affect the result as long as it holds the correct token +// set. RoPE is baked into K once at write time (using the token's own true absolute +// position) and never recomputed from buffer physical position at read time, so +// reordering physical slots doesn't affect that either. +// +// Cost/efficiency: write_kv_slice_sliding() with Circular layout costs O(num_new_tokens) +// per call, always - no O(window_size) data movement ever occurs, even once the window is +// saturated (unlike the LeftAligned layout, which must shift the surviving tail forward +// on every saturated update). This replaces an earlier ~600ms/generate-step bulk-copy +// optimization (still used by LeftAligned buffers) with an approach that needs no data +// movement of existing content at all. +// ============================================================================ + +void LLMBlockKVCacheStrategy::bind_swa_window_views() { + namespace uu = ov::npuw::util; + for (auto& [layer_idx, layer_managers] : m_kv_cache_block_managers) { + if (!layer_managers.is_sliding) { + continue; + } + const std::string layer_idx_str = std::to_string(layer_idx); + + auto bind_one = [&](const ov::SoPtr& window, uint32_t kv_dim, const char* kv_type) { + if (!window) { + return; + } + const uint32_t num_slots = static_cast(window->get_shape()[kv_dim]) / m_block_size; + for (uint32_t block_idx = 0; block_idx < num_slots; ++block_idx) { + const std::string input_name = make_numbered_block_input_name(kv_type, layer_idx_str, block_idx); + const uint32_t start = block_idx * m_block_size; + const uint32_t end = start + m_block_size; + auto view = uu::make_tensor_slice(window, kv_dim, start, end); + + auto prefill_it = m_prefill_classified_in_ports.find(input_name); + if (prefill_it != m_prefill_classified_in_ports.end()) { + m_req.m_prefill_request->set_tensor(prefill_it->second.port, view); + } + for (const auto& generate_request : m_req.m_generate_requests) { + const auto& classified = m_gen_classified_in_ports.at(generate_request); + auto gen_it = classified.find(input_name); + if (gen_it != classified.end()) { + generate_request->set_tensor(gen_it->second.port, view); + } + } + } + LOG_VERB("[SWA] Layer " << layer_idx << " " << kv_type << ": bound " << num_slots + << " numbered port(s) as static views into the window buffer"); + }; + + bind_one(layer_managers.swa_key_window, layer_managers.swa_key_dim, "key"); + bind_one(layer_managers.swa_value_window, layer_managers.swa_value_dim, "value"); + } +} + +void LLMBlockKVCacheStrategy::update_swa_windows_from_prefill(uint32_t current_prompts_len) { + namespace uu = ov::npuw::util; + if (m_kv_cache_block_managers.empty()) { + return; + } + auto& kvcache_desc = m_req.m_npuw_llm_compiled_model->m_kvcache_desc; + const uint32_t tokens_before = kvcache_desc.num_stored_tokens - current_prompts_len; + + for (const auto& [layer_idx, layer_managers] : m_kv_cache_block_managers) { + if (!layer_managers.is_sliding) { + continue; + } + const std::string layer_str = std::to_string(layer_idx); + for (const bool is_key : {true, false}) { + const auto& window = is_key ? layer_managers.swa_key_window : layer_managers.swa_value_window; + const uint32_t dst_kv_dim = is_key ? layer_managers.swa_key_dim : layer_managers.swa_value_dim; + if (!window) { + continue; + } + const std::string output_name = "present." + layer_str + "." + (is_key ? "key" : "value"); + auto port_it = m_req.m_prefill_out_ports.find(output_name); + if (port_it == m_req.m_prefill_out_ports.end()) { + continue; + } + const uint32_t src_kv_dim = (!is_key && kvcache_desc.v_tensors_transposed_pre) ? 3u : kvcache_desc.dim; + auto src_tensor = m_req.m_prefill_request->get_tensor(port_it->second); + uu::write_kv_slice_sliding(window, + src_tensor, + dst_kv_dim, + src_kv_dim, + tokens_before, + current_prompts_len, + uu::SlidingBufferLayout::Circular); + } + } +} + +void LLMBlockKVCacheStrategy::update_swa_windows_generate(uint32_t input_tokens_len) { + namespace uu = ov::npuw::util; + if (m_kv_cache_block_managers.empty()) { + return; + } + auto& kvcache_desc = m_req.m_npuw_llm_compiled_model->m_kvcache_desc; + const uint32_t tokens_before = kvcache_desc.num_stored_tokens - input_tokens_len; + + for (const auto& [layer_idx, layer_managers] : m_kv_cache_block_managers) { + if (!layer_managers.is_sliding) { + continue; + } + const std::string layer_str = std::to_string(layer_idx); + for (const bool is_key : {true, false}) { + const auto& window = is_key ? layer_managers.swa_key_window : layer_managers.swa_value_window; + const uint32_t dst_kv_dim = is_key ? layer_managers.swa_key_dim : layer_managers.swa_value_dim; + if (!window) { + continue; + } + const std::string output_name = "present." + layer_str + "." + (is_key ? "key" : "value"); + auto port_it = m_req.m_kvcache_out_ports.find(output_name); + if (port_it == m_req.m_kvcache_out_ports.end()) { + continue; + } + const uint32_t src_kv_dim = (!is_key && kvcache_desc.v_tensors_transposed_gen) ? 3u : kvcache_desc.dim; + auto src_tensor = m_req.m_kvcache_request->get_tensor(port_it->second); + uu::write_kv_slice_sliding(window, + src_tensor, + dst_kv_dim, + src_kv_dim, + tokens_before, + input_tokens_len, + uu::SlidingBufferLayout::Circular); + } + } } // ============================================================================ @@ -539,6 +776,9 @@ bool LLMBlockKVCacheStrategy::redirect_prefill_outputs_to_new_blocks( size_t total_blocks_allocated = 0; size_t total_outputs_bound = 0; + // Only ever called with a non-null manager, i.e. non-sliding (full-attention) layers - + // SWA layers always have a null key_manager/value_manager (see LayerBlockManagers) and + // are handled separately by update_swa_windows_from_prefill(). auto process_kv_blocks = [&](uint32_t layer_idx, KVCacheBlockManager* manager, const char* kv_type_name) -> bool { uint32_t start_pos = current_position; uint32_t end_pos = current_position + num_new_tokens; @@ -577,11 +817,7 @@ bool LLMBlockKVCacheStrategy::redirect_prefill_outputs_to_new_blocks( auto allocated_blocks = manager->get_allocated_blocks(); while (allocated_blocks.size() < blocks_needed) { auto new_block_id = manager->allocate_block(); - OPENVINO_ASSERT(new_block_id.has_value(), - "Failed to allocate ", - kv_type_name, - " block for layer ", - layer_idx); + OPENVINO_ASSERT(new_block_id.has_value(), "Failed to allocate ", kv_type_name, " block for layer ", layer_idx); total_blocks_allocated++; allocated_blocks = manager->get_allocated_blocks(); } @@ -688,15 +924,21 @@ void LLMBlockKVCacheStrategy::update_generate_bindings(uint32_t old_num_tokens, const uint32_t old_block_idx = helper.get_block_index_for_position(old_num_tokens - 1); const uint32_t final_block_idx = helper.get_block_index_for_position(new_num_tokens - 1); + // This function only ever runs for non-sliding (full-attention) layers - SWA + // layers always have a null key_manager/value_manager (see LayerBlockManagers) + // and are handled separately by update_swa_windows_generate(). So block indices + // map directly to allocated_blocks / numbered_input_ports slots with no eviction + // translation needed. if (final_block_idx > old_block_idx) { // A block boundary was crossed: bind each newly entered block. // In standard decoding (input_tokens_len=1) this loop runs exactly once. // In speculative decoding it may run more times if input_tokens_len > block_size. // - Numbered blocks: zero-copy via set_tensor(). // - Tail block: full copy from token 0 of the block. - for (uint32_t bidx = old_block_idx + 1; - bidx <= final_block_idx && bidx < static_cast(allocated_blocks.size()); - ++bidx) { + for (uint32_t bidx = old_block_idx + 1; bidx <= final_block_idx; ++bidx) { + if (bidx >= allocated_blocks.size()) { + continue; + } const uint32_t block_id = allocated_blocks[bidx]; if (helper.should_treat_as_tail(bidx)) { copy_block_to_tail_input(manager->get_block_tensor(block_id), @@ -736,11 +978,7 @@ void LLMBlockKVCacheStrategy::update_generate_bindings(uint32_t old_num_tokens, for (const auto& [layer_idx, layer_managers] : m_kv_cache_block_managers) { const auto& layer_helpers = variant_layer_helpers.at(layer_idx); if (layer_managers.key_manager) { - update_blocks(layer_idx, - layer_managers.key_manager.get(), - layer_helpers.key_helper, - kvcache_desc.dim, - "key"); + update_blocks(layer_idx, layer_managers.key_manager.get(), layer_helpers.key_helper, kvcache_desc.dim, "key"); } if (layer_managers.value_manager) { const uint32_t kv_dim = kvcache_desc.v_tensors_transposed_gen ? 3u : kvcache_desc.dim; @@ -780,6 +1018,12 @@ void LLMBlockKVCacheStrategy::copy_outputs_to_blocks(const std::shared_ptrsecond; auto& block_manager = is_key ? layer_managers.key_manager : layer_managers.value_manager; + if (!block_manager) { + // SWA layers never have a key_manager/value_manager (see LayerBlockManagers) - + // they are handled separately by update_swa_windows_from_prefill()/ + // update_swa_windows_generate() via write_kv_slice_sliding(). + continue; + } const uint32_t kv_dim = (!is_key && v_transposed) ? 3u : kvcache_desc.dim; auto src_tensor = request->get_tensor(src_ports.at(output_name)); @@ -796,13 +1040,16 @@ void LLMBlockKVCacheStrategy::copy_outputs_to_blocks(const std::shared_ptrget_allocated_blocks(); - for (uint32_t b = static_cast(allocated_blocks.size()); b <= end_block_idx; ++b) { - OPENVINO_ASSERT(block_manager->allocate_block().has_value(), - "Failed to allocate block for KV cache — pool exhausted"); + // Allocate any blocks that do not yet exist (non-sliding / full-attention layers + // only - SWA layers are filtered out above). Grows until kvcache_desc.total_size. + { + auto allocated_blocks = block_manager->get_allocated_blocks(); + for (uint32_t b = static_cast(allocated_blocks.size()); b <= end_block_idx; ++b) { + OPENVINO_ASSERT(block_manager->allocate_block().has_value(), + "Failed to allocate block for KV cache — pool exhausted"); + } } - allocated_blocks = block_manager->get_allocated_blocks(); + auto allocated_blocks = block_manager->get_allocated_blocks(); // Write tokens across blocks. // The first block may be partially filled: start writing at (start_pos % block_size). @@ -834,7 +1081,29 @@ void LLMBlockKVCacheStrategy::create_block_managers_and_helpers() { m_block_size = block_size; const uint32_t max_blocks = (compiled_model->m_kvcache_desc.total_size + block_size - 1) / block_size; - LOG_INFO("Block configuration: size=" << block_size << " tokens, max_blocks=" << max_blocks); + // Sliding-window (SWA) layers get a much smaller, capped block pool sized to the + // attention window instead of the full kvcache_desc.total_size: once the pool fills + // up, KVCacheBlockManager::ensure_blocks_up_to() evicts the single oldest resident + // block to make room for the newest one, so exactly window_size/block_size blocks + // (the most recent ones) stay resident at all times. + // + // FIXME: window_size not being a multiple of block_size is not yet supported - the + // last partial window block would need its own tail-style handling, similar to the + // tail port used for a non-block-aligned total_size. + uint32_t swa_max_blocks = 0; + if (compiled_model->m_swa_window_size > 0) { + OPENVINO_ASSERT(compiled_model->m_swa_window_size % block_size == 0, + "NPUW block KV cache: SWA window_size (", + compiled_model->m_swa_window_size, + ") must be a multiple of block_size (", + block_size, + "). Non-multiple window sizes are not yet supported by the block-based KV cache " + "strategy."); + swa_max_blocks = compiled_model->m_swa_window_size / block_size; + } + + LOG_INFO("Block configuration: size=" << block_size << " tokens, max_blocks=" << max_blocks + << ", swa_max_blocks=" << swa_max_blocks); // ------------------------------------------------------------------------- // Phase 1: Single scan — discover which layers have key/value block ports. @@ -899,6 +1168,23 @@ void LLMBlockKVCacheStrategy::create_block_managers_and_helpers() { "SplitKVCacheIntoBlocks transformation may be broken."); LayerBlockManagers layer_managers; + layer_managers.is_sliding = compiled_model->is_swa_layer(layer_idx); + + // Detects which dim (2 or 3) of `shape` holds the sequence axis (whichever equals + // block_size), mirroring KVCacheBlockManager's own detection logic, and returns a + // same-rank shape with that dim replaced by `new_len`. + auto make_window_shape = [&](const ov::Shape& shape, uint32_t new_len, uint32_t& detected_dim) { + OPENVINO_ASSERT(shape.size() == 4 && (shape[2] == block_size || shape[3] == block_size), + "NPUW block KV cache: SWA base_shape ", + shape, + " does not have block_size=", + block_size, + " in sequence dimension (expected at dim 2 or 3)"); + detected_dim = (shape[2] == block_size) ? 2u : 3u; + ov::Shape window_shape = shape; + window_shape[detected_dim] = new_len; + return window_shape; + }; if (presence.has_key_numbered_block) { // Invariant B: numbered blocks must start at block_0 @@ -909,12 +1195,24 @@ void LLMBlockKVCacheStrategy::create_block_managers_and_helpers() { " has key blocks but no key_block_0. " "SplitKVCacheIntoBlocks transformation may be broken."); auto first_key_port = m_prefill_classified_in_ports.at(key_block0_name).port; - layer_managers.key_manager = std::make_unique(block_size, - max_blocks, - first_key_port.get_shape(), - first_key_port.get_element_type(), - m_req.m_pre_alloc_device, - compiled_model->get_plugin()); + if (layer_managers.is_sliding) { + // Single persistent window buffer, statically bound to all numbered ports + // as adjacent views (see bind_swa_window_views()) - no KVCacheBlockManager, + // no eviction/rotation. Updated in-place each step via write_kv_slice_sliding(). + const auto window_shape = + make_window_shape(first_key_port.get_shape(), swa_max_blocks * block_size, layer_managers.swa_key_dim); + layer_managers.swa_key_window = ov::npuw::util::allocMem(first_key_port.get_element_type(), + window_shape, + m_req.m_pre_alloc_device, + compiled_model->get_plugin()); + } else { + layer_managers.key_manager = std::make_unique(block_size, + max_blocks, + first_key_port.get_shape(), + first_key_port.get_element_type(), + m_req.m_pre_alloc_device, + compiled_model->get_plugin()); + } } if (presence.has_value_numbered_block) { const std::string value_block0_name = make_numbered_block_input_name("value", layer_idx_str, 0); @@ -924,12 +1222,28 @@ void LLMBlockKVCacheStrategy::create_block_managers_and_helpers() { " has value blocks but no value_block_0. " "SplitKVCacheIntoBlocks transformation may be broken."); auto first_value_port = m_prefill_classified_in_ports.at(value_block0_name).port; - layer_managers.value_manager = std::make_unique(block_size, - max_blocks, - first_value_port.get_shape(), - first_value_port.get_element_type(), - m_req.m_pre_alloc_device, - compiled_model->get_plugin()); + if (layer_managers.is_sliding) { + const auto window_shape = make_window_shape(first_value_port.get_shape(), + swa_max_blocks * block_size, + layer_managers.swa_value_dim); + layer_managers.swa_value_window = ov::npuw::util::allocMem(first_value_port.get_element_type(), + window_shape, + m_req.m_pre_alloc_device, + compiled_model->get_plugin()); + } else { + layer_managers.value_manager = std::make_unique(block_size, + max_blocks, + first_value_port.get_shape(), + first_value_port.get_element_type(), + m_req.m_pre_alloc_device, + compiled_model->get_plugin()); + } + } + + if (layer_managers.is_sliding) { + LOG_INFO("[SWA] Layer " << layer_idx << ": sliding, window buffer capacity=" << swa_max_blocks * block_size + << " tokens (window_size=" << compiled_model->m_swa_window_size + << ", block_size=" << block_size << ")"); } m_kv_cache_block_managers[layer_idx] = std::move(layer_managers); @@ -940,6 +1254,31 @@ void LLMBlockKVCacheStrategy::create_block_managers_and_helpers() { for (const auto& [generate_request, classified] : m_gen_classified_in_ports) { m_variant_block_binding_helpers[generate_request] = build_variant_layer_helpers(classified, block_size); } + + // Diagnostic: for every SWA layer, log how many numbered/tail input ports the + // compiled model actually exposes vs. the block pool capacity we just computed. + // These two must match (swa_max_blocks == numbered_input_ports.size(), no tail port + // when window_size % block_size == 0) - a mismatch here would mean some numbered + // ports never get bound to real data, which is a likely source of wrong results. + if (!m_variant_block_binding_helpers.empty()) { + const auto& first_variant_helpers = m_variant_block_binding_helpers.begin()->second; + for (const auto& [layer_idx, layer_managers] : m_kv_cache_block_managers) { + if (!layer_managers.is_sliding) { + continue; + } + auto it = first_variant_helpers.find(layer_idx); + if (it == first_variant_helpers.end()) { + LOG_WARN("[SWA] Layer " << layer_idx << ": sliding, but no generate binding helpers found!"); + continue; + } + const auto& kh = it->second.key_helper; + const auto& vh = it->second.value_helper; + LOG_INFO("[SWA] Layer " << layer_idx << ": generate model ports — key: " << kh.numbered_input_ports.size() + << " numbered" << (kh.has_tail_input() ? " + tail" : " (no tail)") + << ", value: " << vh.numbered_input_ports.size() << " numbered" + << (vh.has_tail_input() ? " + tail" : " (no tail)")); + } + } } } // namespace npuw diff --git a/src/plugins/intel_npu/src/plugin/npuw/llm_block_kvcache_strategy.hpp b/src/plugins/intel_npu/src/plugin/npuw/llm_block_kvcache_strategy.hpp index ed6f396ea8ad75..246b20478a6cd4 100644 --- a/src/plugins/intel_npu/src/plugin/npuw/llm_block_kvcache_strategy.hpp +++ b/src/plugins/intel_npu/src/plugin/npuw/llm_block_kvcache_strategy.hpp @@ -32,8 +32,46 @@ enum class BlockParamKind { /// @brief Pair of key/value block managers for one transformer layer. struct LayerBlockManagers { + // Used for NON-sliding (full-attention) layers only: a growable pool of block_size-wide + // blocks, zero-copy bound directly to the model's numbered block ports. std::unique_ptr key_manager; std::unique_ptr value_manager; + + // True for a sliding-window-attention (SWA) layer. SWA layers do NOT use key_manager / + // value_manager (both stay null) - a per-token-exact sliding window cannot be + // represented by whole-block eviction/rotation alone (block-granularity eviction can + // only keep/evict in units of block_size, which cannot reproduce an exact window_size + // cutoff whenever it doesn't land on a block boundary - "torn" mid-block positions). + // Instead SWA layers use a single persistent swa_key_window / swa_value_window buffer + // (see below), updated via ov::npuw::util::write_kv_slice_sliding() using + // SlidingBufferLayout::Circular - see the detailed rationale above + // update_swa_windows_from_prefill()/update_swa_windows_generate() in the .cpp file for + // why physical column order within the window doesn't need to be chronological. + bool is_sliding = false; + + // SWA layers only: one contiguous window_size-wide buffer per key/value, allocated once + // in create_block_managers_and_helpers() and never reallocated. All numbered input ports + // for this layer (past_key_values.N.key_block_0..k-1), across the prefill request AND + // every generate variant, are bound ONCE (see bind_swa_window_views()) as adjacent + // block_size-wide VIEWS into this single buffer, and are never rebound again - only the + // buffer's CONTENT changes (via write_kv_slice_sliding()), so the k views always show the + // latest data automatically. + // + // NB: an oversized-buffer + periodic-compaction + "slide the view forward" scheme was + // tried and reverted - the NPU backend's zero-copy remote tensor binding requires each + // port's bound memory to stay at a FIXED (block-aligned, established-once) address; + // rebinding to an arbitrary shifting offset every step triggered "Strided remote tensor + // is not supported for this port!" and produced wrong results. The current + // Circular-layout write scheme keeps every port's binding fixed for its entire lifetime + // and never rebinds anything - only the write OFFSET within the fixed buffer changes. + ov::SoPtr swa_key_window; + ov::SoPtr swa_value_window; + + // The dimension (2 or 3) within swa_key_window/swa_value_window's shape that represents + // the token/sequence axis - mirrors the dim auto-detected by KVCacheBlockManager for + // non-sliding layers (base_shape[2] or base_shape[3] == block_size). + uint32_t swa_key_dim = 0; + uint32_t swa_value_dim = 0; }; /// @brief Helper for block binding: determines whether a block uses zero-copy numbered binding @@ -144,6 +182,23 @@ class LLMBlockKVCacheStrategy final : public LLMKVCacheStrategy { void create_block_managers_and_helpers(); + // Bind every SWA layer's numbered input ports (prefill + all generate variants) as + // static, never-rebound VIEWS into its single swa_key_window/swa_value_window buffer. + // Called once from on_initialize(), and again from on_reset() (which - like any other + // numbered block port - overwrites SWA ports with dummy tensors as a side effect of + // releasing non-sliding blocks, so the static views need re-establishing). + void bind_swa_window_views(); + + // Slide SWA layers' window buffers forward using this prefill chunk's freshly computed + // "present.N.key/value" outputs. Called unconditionally at the end of + // on_prefill_chunk_done(), regardless of whether non-sliding layers used the zero-copy + // or copy path for this chunk (SWA layers are never zero-copy-redirected). + void update_swa_windows_from_prefill(uint32_t current_prompts_len); + + // Slide SWA layers' window buffers forward using this generate step's freshly computed + // "present.N.key/value" outputs. Called unconditionally from on_generate_step_done(). + void update_swa_windows_generate(uint32_t input_tokens_len); + // ------------------------------------------------------------------------- // Core KV copy/bind engine (shared by prefill and generate paths) // ------------------------------------------------------------------------- @@ -206,6 +261,13 @@ class LLMBlockKVCacheStrategy final : public LLMKVCacheStrategy { // Block size in tokens — fixed at on_initialize() time, equal to m_prefill_chunk_size. // Cached here to avoid repeated map lookups into m_kv_cache_block_managers. uint32_t m_block_size = 0; + + // Perf diagnostics: running total/count of update_swa_windows_generate() time, so + // on_generate_step_done()'s per-call LOG_DEBUG can also report a running average + // without needing a separate profiling pass. Reset at the start of every conversation + // (on_reset()) so the numbers reflect a single generation, not the whole process lifetime. + float m_swa_update_total_ms = 0.0f; + uint64_t m_swa_update_calls = 0; }; } // namespace npuw diff --git a/src/plugins/intel_npu/src/plugin/npuw/llm_compiled_model.cpp b/src/plugins/intel_npu/src/plugin/npuw/llm_compiled_model.cpp index 47efb2987c72f6..2b3ebad0efd1c2 100644 --- a/src/plugins/intel_npu/src/plugin/npuw/llm_compiled_model.cpp +++ b/src/plugins/intel_npu/src/plugin/npuw/llm_compiled_model.cpp @@ -3,6 +3,8 @@ // #include "llm_compiled_model.hpp" +#include + #include "embedding/embedding_infer_request.hpp" #include "embedding/encoder_embedding_infer_request.hpp" #include "embedding/prepare_embedding_model.hpp" @@ -20,6 +22,7 @@ #include "npuw_transformations/duplicate_shared_kv_concat.hpp" #include "npuw_transformations/lora_stateful_to_stateless.hpp" #include "npuw_transformations/optimize_value_tensors.hpp" +#include "npuw_transformations/patch_sliding_window_kvcache.hpp" #include "npuw_transformations/patch_sliding_window_mask.hpp" #include "npuw_transformations/remove_token_type_ids.hpp" #include "npuw_transformations/replace_deepstack_scatter_with_add.hpp" @@ -469,6 +472,19 @@ std::map any_copy(const ov::AnyMap& params) { return result; } +// Detect Gemma-4 E2B/E4B by the presence of "per_layer_inputs" model inputs +bool has_per_layer_inputs(const std::shared_ptr& model) { + for (const auto& param : model->get_parameters()) { + if (param->get_friendly_name().find("per_layer_inputs") != std::string::npos) { + LOG_INFO("Detected cross-group KV sharing model (Gemma-4 E2B/E4B): " + "found per_layer_inputs parameter - " + << param->get_friendly_name()); + return true; + } + } + return false; +} + // Detect if the model is a Mixture-of-Experts (MoE) architecture // by checking if any node name matches MoE patterns: layers.*.mlp.router or layers.*.mlp.experts bool is_moe_model(const std::shared_ptr& model) { @@ -574,6 +590,65 @@ std::shared_ptr check_and_cut_lm_head(const std::shared_ptr(m_cfg.get<::intel_npu::NPUW_LLM_SLIDING_WINDOW>()); + m_swa_layer_is_sliding.clear(); + + const std::string layer_types = m_cfg.get<::intel_npu::NPUW_LLM_LAYER_TYPES>(); + if (m_swa_window_size == 0 || layer_types.empty()) { + LOG_DEBUG("[SWA] Sliding Window Attention is disabled (window_size=" << m_swa_window_size + << ", layer_types is " + << (layer_types.empty() ? "empty" : "set") + << ")."); + return; + } + + std::stringstream ss(layer_types); + std::string token; + while (std::getline(ss, token, ',')) { + const auto not_space = [](unsigned char c) { + return !std::isspace(c); + }; + token.erase(token.begin(), std::find_if(token.begin(), token.end(), not_space)); + token.erase(std::find_if(token.rbegin(), token.rend(), not_space).base(), token.end()); + m_swa_layer_is_sliding.push_back(token == "sliding_attention"); + } + + std::string pattern; + pattern.reserve(m_swa_layer_is_sliding.size()); + size_t num_sliding = 0; + for (const bool is_sliding : m_swa_layer_is_sliding) { + pattern.push_back(is_sliding ? 'S' : 'F'); + num_sliding += is_sliding ? 1 : 0; + } + LOG_INFO("[SWA] Sliding Window Attention is ENABLED: window_size=" << m_swa_window_size << ", " + << m_swa_layer_is_sliding.size() + << " layers total, " << num_sliding + << " sliding-window layer(s)."); + LOG_DEBUG("[SWA] Layer pattern (S=sliding, F=full attention): " << pattern); +} + +bool ov::npuw::LLMCompiledModel::is_swa_layer(size_t layer_idx) const { + return m_swa_window_size > 0 && layer_idx < m_swa_layer_is_sliding.size() && m_swa_layer_is_sliding[layer_idx]; +} + +namespace { +// Dumps `model` to an .xml/.bin pair next to the current working directory when +// NPUW_DUMP_FULL is enabled, so users can inspect the effect of +// PatchSlidingWindowKVCache on the compiled sub-model. Reuses the existing +// NPUW_DUMP_FULL flag/convention instead of introducing a new dedicated option. +void dump_swa_model_if_requested(const ::intel_npu::Config& cfg, + const std::shared_ptr& model, + const std::string& tag) { + if (!cfg.get<::intel_npu::NPUW_DUMP_FULL>()) { + return; + } + const std::string path = model->get_friendly_name() + "_swa_" + tag + ".xml"; + ov::save_model(model, path); + LOG_INFO("[SWA] Dumped post-transform model to " << path << " for inspection."); +} +} // namespace + // Apply DEVICE_ROUTED MoE transformations to models std::vector> ov::npuw::LLMCompiledModel::create_generate_model_variants( const std::shared_ptr& generate_model, @@ -636,6 +711,19 @@ std::vector> ov::npuw::LLMCompiledModel::create_gener ov::npuw::ReshapeToStatic(max_generation_token_len, kv_size, axes, m_max_lora_rank, whisper_lhs_seq_size) .run_on_model(generate_variant); + if (m_swa_window_size > 0) { + LOG_DEBUG("[SWA] Applying sliding-window KV-cache reduction to generate variant (kv_size=" << kv_size + << ")."); + ov::npuw::PatchSlidingWindowKVCache(m_swa_window_size, + m_swa_layer_is_sliding, + kv_size, + max_generation_token_len, + axes, + /*trim_attention_mask=*/true) + .run_on_model(generate_variant); + dump_swa_model_if_requested(m_cfg, generate_variant, "generate_kv" + std::to_string(kv_size)); + } + // Set unique name for this variant generate_variant->set_friendly_name(generate_model->get_friendly_name() + "_kv" + std::to_string(kv_size)); generate_model_variants.push_back(generate_variant); @@ -824,6 +912,8 @@ ov::npuw::LLMCompiledModel::LLMCompiledModel(const std::shared_ptr& m const uint32_t seq_len_dim = m_cfg.get<::intel_npu::NPUW_LLM_SEQ_LEN_DIM>(); KVAxesPosition axes{batch_dim, seq_len_dim}; + parse_swa_config(); + LOG_DEBUG("Creating kvcache model as clone of passed one."); auto kvcache_model = model->clone(); @@ -902,6 +992,13 @@ ov::npuw::LLMCompiledModel::LLMCompiledModel(const std::shared_ptr& m ov::npuw::DetectAttentionMask detect_mask; detect_mask.run_on_model(kvcache_model); const auto mask_info = detect_mask.get_mask_info(); + { + const char* mask_type_str = mask_info.mask_type == ov::npuw::MaskInfo::MaskType::Causal ? "Causal" + : mask_info.mask_type == ov::npuw::MaskInfo::MaskType::SlidingWindow ? "SlidingWindow" + : "Unknown"; + LOG_INFO("[HFA] DetectAttentionMask result: mask_type=" << mask_type_str + << ", window_size=" << mask_info.window_size); + } if (!m_is_whisper) { LOG_DEBUG("Try patch sliding window attention mask (Phi-3, Gemma-2, Gemma-3, Gemma-4), if it exists."); @@ -968,6 +1065,23 @@ ov::npuw::LLMCompiledModel::LLMCompiledModel(const std::shared_ptr& m true) .run_on_model(prefill_model); } + if (m_swa_window_size > 0) { + // NB: prefill's own attention mask must stay full-width (it is already correctly + // banded per-query-position by PatchSlidingWindowMask above) - only shrink the + // past_key_values buffers here (a no-op for non-chunked prefill, where the past + // buffer is always empty), never trim the mask. + const uint32_t prefill_input_size = m_use_chunk_prefill ? static_cast(m_prefill_chunk_size) + : m_kvcache_desc.max_prompt_size; + LOG_DEBUG("[SWA] Applying sliding-window KV-cache reduction to prefill model."); + ov::npuw::PatchSlidingWindowKVCache(m_swa_window_size, + m_swa_layer_is_sliding, + m_kvcache_desc.max_prompt_size, + prefill_input_size, + axes, + /*trim_attention_mask=*/false) + .run_on_model(prefill_model); + dump_swa_model_if_requested(m_cfg, prefill_model, "prefill"); + } LOG_DEBUG("Make kvcache model with static shapes"); // In case of Gemma3, we should remove `token_type_ids` from generate version of the model, @@ -1051,6 +1165,17 @@ ov::npuw::LLMCompiledModel::LLMCompiledModel(const std::shared_ptr& 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"); @@ -1082,12 +1207,24 @@ ov::npuw::LLMCompiledModel::LLMCompiledModel(const std::shared_ptr& m prefill_config_opt.value_or(get_default_prefill_config(prefill_model, npudesc)).as(); if (prefill_attn_hfa) { - prefill_config[ov::intel_npu::npuw::partitioning::attn_hfa_mask_skipping.name()] = + const bool mask_skipping_enabled = mask_info.mask_type == ov::npuw::MaskInfo::MaskType::Causal || - (mask_info.mask_type == ov::npuw::MaskInfo::MaskType::SlidingWindow && - mask_info.window_size >= max_prompt_len) - ? "YES" - : "NO"; + (mask_info.mask_type == ov::npuw::MaskInfo::MaskType::SlidingWindow && + mask_info.window_size >= max_prompt_len); + prefill_config[ov::intel_npu::npuw::partitioning::attn_hfa_mask_skipping.name()] = + mask_skipping_enabled ? "YES" : "NO"; + LOG_INFO("[HFA] NPUW_ATTN_HFA_MASK_SKIPPING=" << (mask_skipping_enabled ? "YES" : "NO") + << " (mask_type=" + << (mask_info.mask_type == ov::npuw::MaskInfo::MaskType::Causal + ? "Causal" + : mask_info.mask_type == + ov::npuw::MaskInfo::MaskType::SlidingWindow + ? "SlidingWindow" + : "Unknown") + << ", window_size=" << mask_info.window_size + << ", max_prompt_len=" << max_prompt_len + << "). When YES, HFA regular (non-final) prefill tiles " + "skip the explicit mask."); } // NB: GENERATE_HINT is only applicable for default generate config! @@ -1139,6 +1276,19 @@ ov::npuw::LLMCompiledModel::LLMCompiledModel(const std::shared_ptr& m if (m_use_chunk_prefill && (prefill_attn_pyramid || prefill_attn_hfa || prefill_attn_dyn)) { prefill_config["NPUW_ATTN"] = ::intel_npu::NPUW_LLM_PREFILL_ATTENTION_HINT::toString(prefill_attn_hint); merge_config_with(prefill_config, dyn_attn_opts); + // Gemma-4 E2B/E4B: the SWA<->Global boundary subgraphs (FFN tail + Global Q-proj + // prefix) appear in two structurally distinct variants depending on the K/V source + // of the preceding SWA layer: + // Variant A (L3->L4, L8->L9, L13->L14): SWA layer uses its own group K/V -> 3 instances + // Variant B (L18->L19, L23->L24, L28->L29, L33->L34): SWA layer uses L13 + // borrowed K/V (cloned by DuplicateSharedKVConcat) -> 4 instances + // Both variants fall below the default keep_blocks=5 threshold and remain as + // separate FCE compile units. Lower to 3 so both are folded into REP and + // reused across their respective instances. + if (has_per_layer_inputs(prefill_model)) { + prefill_config["NPUW_ONLINE_KEEP_BLOCKS"] = "3"; + LOG_INFO("Gemma-4 cross-group KV model: setting NPUW_ONLINE_KEEP_BLOCKS=3 for prefill"); + } } if (generate_attn_pyramid || generate_attn_hfa || generate_attn_dyn) { @@ -1660,6 +1810,9 @@ std::shared_ptr ov::npuw::LLMCompiledModel::deserial compiled->implement_properties(); // Not serialized. Recomputed from the deserialized config so older blobs stay loadable. compiled->m_enable_continuous_prefill = compiled->m_cfg.get<::intel_npu::NPUW_LLM_ENABLE_CONTINUOUS_PREFILL>(); + // NB: SWA state is fully derived from cached config options, not separately + // serialized - recompute it now that m_cfg is restored. + compiled->parse_swa_config(); // Deserialize KV cache model variants stream & compiled->m_kvcache_sizes; diff --git a/src/plugins/intel_npu/src/plugin/npuw/llm_compiled_model.hpp b/src/plugins/intel_npu/src/plugin/npuw/llm_compiled_model.hpp index a928d4a41b1b0e..eb1f98e852bd84 100644 --- a/src/plugins/intel_npu/src/plugin/npuw/llm_compiled_model.hpp +++ b/src/plugins/intel_npu/src/plugin/npuw/llm_compiled_model.hpp @@ -159,6 +159,21 @@ class LLMCompiledModel : public ov::npuw::ICompiledModel { // routed to the dedicated KV/RoPE-free encoder embedding path. bool m_is_encoder_embedding = false; + // Sliding Window Attention (SWA) support: per-layer KV-cache window capping for + // hybrid sliding/full-attention models (e.g. Gemma4). Both fields are derived from + // the cached NPUW_LLM_SLIDING_WINDOW / NPUW_LLM_LAYER_TYPES config options - they are + // NOT separately serialized, they are recomputed by parse_swa_config() both in the + // constructor and right after m_cfg is restored in deserialize(). + uint32_t m_swa_window_size = 0; // 0 == Sliding Window Attention disabled + std::vector m_swa_layer_is_sliding; // per-layer flag, indexed by decoder layer id + + // Parses NPUW_LLM_SLIDING_WINDOW / NPUW_LLM_LAYER_TYPES from m_cfg into + // m_swa_window_size / m_swa_layer_is_sliding. + void parse_swa_config(); + + // True if SWA is enabled and layer_idx is configured as a sliding-window layer. + bool is_swa_layer(size_t layer_idx) const; + // Create generate model variants with different sizes std::vector> create_generate_model_variants( const std::shared_ptr& generate_model, diff --git a/src/plugins/intel_npu/src/plugin/npuw/llm_continuous_kvcache_strategy.cpp b/src/plugins/intel_npu/src/plugin/npuw/llm_continuous_kvcache_strategy.cpp index 4b54a4694ddec2..afab9db9fb6f56 100644 --- a/src/plugins/intel_npu/src/plugin/npuw/llm_continuous_kvcache_strategy.cpp +++ b/src/plugins/intel_npu/src/plugin/npuw/llm_continuous_kvcache_strategy.cpp @@ -4,6 +4,8 @@ #include "llm_continuous_kvcache_strategy.hpp" +#include + #include "infer_request_utils.hpp" #include "llm_infer_request.hpp" #include "util.hpp" @@ -127,8 +129,18 @@ void LLMContinuousKVCacheStrategy::on_generate_variant_switch(const std::shared_ ? 3u : kvcache_desc.dim; - auto src_slice = uu::make_tensor_slice(src, kv_dim, 0u, num_stored); - auto dst_slice = uu::make_tensor_slice(dst, kv_dim, 0u, num_stored); + // For SWA layers, the past-KV capacity (dim size of `src`/`dst`) may differ between + // generate variants (it is derived as window_size - variant's own kv_size), and may + // be smaller than `num_stored`. Buffers are left-aligned holding their most recent + // valid() tokens in chronological order, so when clamping we must keep the TAIL + // (most recent) `copy_count` tokens of src's valid region, not its head. + const uint32_t src_cap = static_cast(src->get_shape()[kv_dim]); + const uint32_t dst_cap = static_cast(dst->get_shape()[kv_dim]); + const uint32_t src_valid = std::min(num_stored, src_cap); + const uint32_t copy_count = std::min(src_valid, dst_cap); + + auto src_slice = uu::make_tensor_slice(src, kv_dim, src_valid - copy_count, src_valid); + auto dst_slice = uu::make_tensor_slice(dst, kv_dim, 0u, copy_count); // Copy via a temporary CPU buffer to avoid aliasing (src and dst share backing memory). auto tmp = uu::allocMem(src->get_element_type(), src_slice->get_shape(), "CPU", nullptr); @@ -136,8 +148,6 @@ void LLMContinuousKVCacheStrategy::on_generate_variant_switch(const std::shared_ uu::copy_tensor_by_dim(tmp, dst_slice, kv_dim, kv_dim); } } - -// on_generate_step_done: persist the new token's KV output into the past KV input buffer // so the next generate step sees the updated context. void LLMContinuousKVCacheStrategy::on_generate_step_done(uint32_t input_tokens_len) { const bool v_transposed = m_req.m_npuw_llm_compiled_model->m_kvcache_desc.v_tensors_transposed_gen; diff --git a/src/plugins/intel_npu/src/plugin/npuw/llm_infer_base_request.cpp b/src/plugins/intel_npu/src/plugin/npuw/llm_infer_base_request.cpp index e74ceb9236da66..670e1214b36340 100644 --- a/src/plugins/intel_npu/src/plugin/npuw/llm_infer_base_request.cpp +++ b/src/plugins/intel_npu/src/plugin/npuw/llm_infer_base_request.cpp @@ -30,22 +30,21 @@ void ov::npuw::LLMInferBaseRequest::update_kvcache_for( } auto dst_tensor = request->get_tensor(in_ports.at(input_name)); const auto& kv_dim = (output_name.find("value") != std::string::npos && v_transposed) ? 3u : kvcache_desc.dim; - auto dst_slice = uu::make_tensor_slice(dst_tensor, - kv_dim, - kvcache_desc.num_stored_tokens - num_tokens, - kvcache_desc.num_stored_tokens); auto src_tensor = request->get_tensor(out_ports.at(output_name)); // NOTE: Sometimes present kv layer can contain greater seq_len // than was sent to be processed uint32_t src_seq_len = static_cast(src_tensor->get_shape()[kv_dim]); OPENVINO_ASSERT(num_tokens <= src_seq_len); - if (src_seq_len > num_tokens) { - auto src_slice = uu::make_tensor_slice(src_tensor, kv_dim, src_seq_len - num_tokens, src_seq_len); - uu::copy_tensor_by_dim(src_slice, dst_slice, kv_dim, kv_dim); - } else { - uu::copy_tensor_by_dim(src_tensor, dst_slice, kv_dim, kv_dim); - } + // write_kv_slice_sliding() clamps against dst_tensor's own capacity, so this is + // also correct (and a no-op change in behavior) for non-SWA layers, where + // capacity always covers the whole context. + uu::write_kv_slice_sliding(dst_tensor, + src_tensor, + kv_dim, + kv_dim, + kvcache_desc.num_stored_tokens - num_tokens, + num_tokens); } } diff --git a/src/plugins/intel_npu/src/plugin/npuw/llm_infer_request.cpp b/src/plugins/intel_npu/src/plugin/npuw/llm_infer_request.cpp index 2af2e2941c5f83..d20dd057921aef 100644 --- a/src/plugins/intel_npu/src/plugin/npuw/llm_infer_request.cpp +++ b/src/plugins/intel_npu/src/plugin/npuw/llm_infer_request.cpp @@ -650,6 +650,18 @@ void ov::npuw::LLMInferRequest::copy_kvcache() { // Create backup of past KV tensor when buffer sharing is enabled to prevent data corruption // This is necessary because subsequent copy operations would overwrite the shared buffer auto prefill_past_kv = m_prefill_request->get_tensor(m_prefill_in_ports.at(input_name)); + + // For SWA layers, prefill_past_kv's OWN capacity (its shape at pre_kv_dim) is shrunk by + // PatchSlidingWindowKVCache to window_size, which can be SMALLER than the logical + // tokens_in_past_chunks count once the prompt spans enough chunks to saturate the window + // - slicing to the raw (unclamped) tokens_in_past_chunks would then read out of bounds. + // Per the KV buffer invariant, the tensor always holds exactly its min(total, capacity) + // MOST RECENT tokens, LEFT-aligned - so the valid prefix to read is [0, valid_past_chunks), + // which is already exactly that "most recent" prefix by construction (no extra math needed). + const auto pre_capacity = static_cast(prefill_past_kv->get_shape()[pre_kv_dim]); + const uint32_t valid_past_chunks = + std::min(static_cast(tokens_in_past_chunks), pre_capacity); + ov::SoPtr tmp_dense_kv_tensor; ov::SoPtr prefill_past_kv_chunks; if (m_past_kv_bound) { @@ -658,49 +670,59 @@ void ov::npuw::LLMInferRequest::copy_kvcache() { m_pre_alloc_device, m_npuw_llm_compiled_model->get_plugin()); prefill_past_kv->copy_to(tmp_dense_kv_tensor._ptr); - prefill_past_kv_chunks = make_tensor_slice(tmp_dense_kv_tensor, - pre_kv_dim, - 0u, - static_cast(tokens_in_past_chunks)); + prefill_past_kv_chunks = make_tensor_slice(tmp_dense_kv_tensor, pre_kv_dim, 0u, valid_past_chunks); } else { - prefill_past_kv_chunks = make_tensor_slice(prefill_past_kv, - pre_kv_dim, - 0u, - static_cast(tokens_in_past_chunks)); + prefill_past_kv_chunks = make_tensor_slice(prefill_past_kv, pre_kv_dim, 0u, valid_past_chunks); } - auto kvcache_past_kv_chunks = uu::make_tensor_slice(kvcache_in_tensor, - gen_kv_dim, - 0u, - static_cast(tokens_in_past_chunks)); - - uu::copy_tensor_by_dim(prefill_past_kv_chunks, kvcache_past_kv_chunks, pre_kv_dim, gen_kv_dim); + // prefill_past_kv_chunks may hold FEWER than tokens_in_past_chunks tokens for a saturated + // SWA layer (see above) - pass the full LOGICAL tokens_in_past_chunks as num_new_tokens + // regardless: write_kv_slice_sliding() re-clamps against the source tensor's own (possibly + // smaller) length internally and extracts its TAIL (most recent) portion, which is exactly + // correct since prefill_past_kv_chunks is itself already the most-recent-valid_past_chunks + // prefix. + uu::write_kv_slice_sliding(kvcache_in_tensor, + prefill_past_kv_chunks, + gen_kv_dim, + pre_kv_dim, + 0u, + static_cast(tokens_in_past_chunks)); } - // Copy part 2 KV results + // Copy part 2 KV results. prefill_out_tensor here is the LAST chunk's own "present" OUTPUT - + // a pure function of this forward call's own input length (prefill_chunk_size), computed fresh + // each call and NEVER shrunk by PatchSlidingWindowKVCache (only the "past" INPUT Parameter is + // shrunk for SWA layers) - so this slice bound is safe for SWA and non-SWA layers alike. auto prefill_present_kv_chunk = uu::make_tensor_slice(prefill_out_tensor, pre_kv_dim, static_cast(prefill_chunk_size - m_tokens_in_present_chunk), static_cast(prefill_chunk_size)); - auto kvcache_last_kv_chunk = uu::make_tensor_slice(kvcache_in_tensor, - gen_kv_dim, - static_cast(tokens_in_past_chunks), - kvcache_desc.num_stored_tokens); - - uu::copy_tensor_by_dim(prefill_present_kv_chunk, kvcache_last_kv_chunk, pre_kv_dim, gen_kv_dim); + // num_stored_tokens_before is the logical token count represented by part 1 + // above; write_kv_slice_sliding re-derives the capacity-clamped valid count + // from it internally, consistently with what part 1 actually wrote. + uu::write_kv_slice_sliding(kvcache_in_tensor, + prefill_present_kv_chunk, + gen_kv_dim, + pre_kv_dim, + static_cast(tokens_in_past_chunks), + static_cast(m_tokens_in_present_chunk)); } else { - auto prefill_out_slice = - uu::make_tensor_slice(prefill_out_tensor, + // prefill_out_tensor follows the "present/output" convention (right-aligned + // valid tail of length num_stored_tokens within the max_prompt_size buffer); + // this is the whole (non-chunked) prefill's own single-shot "present" output, a pure + // function of the model's own input_size (== max_prompt_size here) - like chunked + // prefill's part 2, it is NEVER shrunk by PatchSlidingWindowKVCache for SWA layers (only + // the "past" INPUT Parameter is), so no pre-slice/clamp is needed here: the tensor's own + // full shape at pre_kv_dim is always >= num_stored_tokens. write_kv_slice_sliding extracts + // the correct (capacity-clamped) tail itself. + uu::write_kv_slice_sliding(kvcache_in_tensor, + prefill_out_tensor, + gen_kv_dim, pre_kv_dim, - kvcache_desc.max_prompt_size - kvcache_desc.num_stored_tokens, - kvcache_desc.max_prompt_size); - - auto kvcache_in_slice = - uu::make_tensor_slice(kvcache_in_tensor, gen_kv_dim, 0u, kvcache_desc.num_stored_tokens); - - uu::copy_tensor_by_dim(prefill_out_slice, kvcache_in_slice, pre_kv_dim, gen_kv_dim); + 0u, + kvcache_desc.num_stored_tokens); } }); LOG_DEBUG("Done."); @@ -729,22 +751,21 @@ void ov::npuw::LLMInferRequest::update_kvcache_for( auto dst_tensor = request->get_tensor(in_ports.at(input_name)); const auto& kv_dim = (output_name.find("value") != std::string::npos && v_transposed) ? 3u : kvcache_desc.dim; - auto dst_slice = uu::make_tensor_slice(dst_tensor, - kv_dim, - kvcache_desc.num_stored_tokens - num_tokens, - kvcache_desc.num_stored_tokens); auto src_tensor = request->get_tensor(out_ports.at(output_name)); // NOTE: Sometimes present kv layer can contain greater seq_len // than was sent to be processed uint32_t src_seq_len = static_cast(src_tensor->get_shape()[kv_dim]); OPENVINO_ASSERT(num_tokens <= src_seq_len); - if (src_seq_len > num_tokens) { - auto src_slice = uu::make_tensor_slice(src_tensor, kv_dim, src_seq_len - num_tokens, src_seq_len); - uu::copy_tensor_by_dim(src_slice, dst_slice, kv_dim, kv_dim); - } else { - uu::copy_tensor_by_dim(src_tensor, dst_slice, kv_dim, kv_dim); - } + // write_kv_slice_sliding() clamps against dst_tensor's own capacity, so this is + // also correct (and a no-op change in behavior) for non-SWA layers, where + // capacity always covers the whole context. + uu::write_kv_slice_sliding(dst_tensor, + src_tensor, + kv_dim, + kv_dim, + kvcache_desc.num_stored_tokens - num_tokens, + num_tokens); } } diff --git a/src/plugins/intel_npu/src/plugin/npuw/npuw_transformations/detect_causal_mask.cpp b/src/plugins/intel_npu/src/plugin/npuw/npuw_transformations/detect_causal_mask.cpp index abc15860d43a3b..2e3fff58c6d60d 100644 --- a/src/plugins/intel_npu/src/plugin/npuw/npuw_transformations/detect_causal_mask.cpp +++ b/src/plugins/intel_npu/src/plugin/npuw/npuw_transformations/detect_causal_mask.cpp @@ -5,7 +5,10 @@ #include "detect_causal_mask.hpp" #include +#include +#include +#include "../util.hpp" #include "openvino/op/ops.hpp" #include "openvino/op/scaled_dot_product_attention.hpp" #include "openvino/pass/graph_rewrite.hpp" @@ -221,6 +224,55 @@ class DefaultSWAMatcher final : public ov::pass::MatcherPass { } }; +// ============================================================================ +// Matches the "masked_fill"-style Gemma-4 sliding-window mask (the newer +// export variant handled by Gemma4MaskedFillSlidingMaskMatcher in +// sliding_window_mask.cpp). The sliding-window check here is a single +// GreaterEqual combined via Select/masked_fill instead of a separate +// LessEqual+Greater pair combined via BitwiseAnd/LogicalAnd/BitwiseOr like +// the matchers above: +// +// col_pos = Unsqueeze(Range(...)) +// row_pos = Add(Unsqueeze(Range(...)), any) +// beyond_window = GreaterEqual(Subtract(row_pos, col_pos), window_size) +// sliding_mask = Select(Unsqueeze(Unsqueeze(beyond_window)), any, causal_mask) +// +// None of the matchers above recognize this shape. Left unmatched, the +// standalone `causal_mask` operand of the Select (still built via a plain +// LessEqual(range_chain, range_chain), which StandardCausalMatcher *does* +// recognize) is the only thing that fires, mis-tagging the whole model's +// mask as plain Causal and wrongly enabling HFA mask-skipping for the SWA +// layers (silently reverting sliding-window prefill to full causal +// attention once the context grows past the window). +// ============================================================================ +class MaskedFillSlidingMatcher final : public ov::pass::MatcherPass { +public: + OPENVINO_MATCHER_PASS_RTTI("ov::npuw::MaskedFillSlidingMatcher"); + explicit MaskedFillSlidingMatcher(ov::npuw::MaskInfo& mask_info) { + auto key_range = opp::wrap_type({opp::any_input(), opp::any_input(), opp::any_input()}); + auto col_pos = opp::wrap_type({key_range, opp::any_input()}); + + auto local_arange = opp::wrap_type({opp::any_input(), opp::any_input(), opp::any_input()}); + auto local_arange_row = opp::wrap_type({local_arange, opp::any_input()}); + auto row_pos = opp::wrap_type({local_arange_row, opp::any_input()}); + + auto row_minus_col = opp::wrap_type({row_pos, col_pos}); + auto window_const = opp::wrap_type(); + auto beyond_window = opp::wrap_type({row_minus_col, window_const}); + auto beyond_window_u1 = opp::wrap_type({beyond_window, opp::any_input()}); + auto beyond_window_u2 = opp::wrap_type({beyond_window_u1, opp::any_input()}); + auto sliding_mask = opp::wrap_type({beyond_window_u2, opp::any_input(), opp::any_input()}); + + auto callback = [=, &mask_info](opp::Matcher& m) { + const int64_t w = get_window_size(m.get_pattern_value_map().at(window_const).get_node_shared_ptr()); + if (w > 0) + mask_info = {ov::npuw::MaskInfo::MaskType::SlidingWindow, w}; + return false; + }; + register_matcher(std::make_shared(sliding_mask, "MaskedFillSliding"), callback); + } +}; + #ifdef __GNUC__ # pragma GCC diagnostic pop #endif @@ -236,6 +288,7 @@ bool DetectAttentionMask::run_on_model(const std::shared_ptr& model) detector.add_matcher(m_mask_info); detector.add_matcher(m_mask_info); detector.add_matcher(m_mask_info); + detector.add_matcher(m_mask_info); detector.add_matcher(m_mask_info); detector.add_matcher(m_mask_info); detector.add_matcher(m_mask_info); @@ -244,4 +297,63 @@ bool DetectAttentionMask::run_on_model(const std::shared_ptr& 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& 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 visited; + std::queue> 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(node) || ov::is_type(node) || + ov::is_type(node)) { + for (size_t i = 0; i < node->get_input_size(); ++i) { + if (ov::is_type(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& 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(mask_type); + m_annotations.push_back({pattern.add_node->get_friendly_name(), mask_type}); + } + return false; +} + } // namespace ov::npuw diff --git a/src/plugins/intel_npu/src/plugin/npuw/npuw_transformations/detect_causal_mask.hpp b/src/plugins/intel_npu/src/plugin/npuw/npuw_transformations/detect_causal_mask.hpp index 283de4ad216059..3e1f653b6f19b1 100644 --- a/src/plugins/intel_npu/src/plugin/npuw/npuw_transformations/detect_causal_mask.hpp +++ b/src/plugins/intel_npu/src/plugin/npuw/npuw_transformations/detect_causal_mask.hpp @@ -4,6 +4,9 @@ #pragma once +#include +#include + #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& model) override; + + // Collected per-SDPA mask types from the most recent run_on_model() call. + const std::vector& get_annotations() const { + return m_annotations; + } + + // Convenience helper: returns only mask types in traversal order. + std::vector get_mask_types() const { + std::vector 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 m_annotations; +}; + } // namespace ov::npuw diff --git a/src/plugins/intel_npu/src/plugin/npuw/npuw_transformations/patch_sliding_window_kvcache.cpp b/src/plugins/intel_npu/src/plugin/npuw/npuw_transformations/patch_sliding_window_kvcache.cpp new file mode 100644 index 00000000000000..c7463178220ecf --- /dev/null +++ b/src/plugins/intel_npu/src/plugin/npuw/npuw_transformations/patch_sliding_window_kvcache.cpp @@ -0,0 +1,617 @@ +// Copyright (C) 2018-2026 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#include "patch_sliding_window_kvcache.hpp" + +#include +#include +#include +#include +#include + +#include "../logging.hpp" +#include "openvino/core/validation_util.hpp" +#include "openvino/op/add.hpp" +#include "openvino/op/broadcast.hpp" +#include "openvino/op/concat.hpp" +#include "openvino/op/constant.hpp" +#include "openvino/op/gather.hpp" +#include "openvino/op/less.hpp" +#include "openvino/op/maximum.hpp" +#include "openvino/op/parameter.hpp" +#include "openvino/op/range.hpp" +#include "openvino/op/reshape.hpp" +#include "openvino/op/scaled_dot_product_attention.hpp" +#include "openvino/op/select.hpp" +#include "openvino/op/slice.hpp" +#include "openvino/op/subtract.hpp" + +namespace { + +// Matches e.g. "past_key_values.3.key" / "past_key_values.12.value" -> layer index 3 / 12. +const std::regex& kv_param_regex() { + static const std::regex re(R"(past_key_values\.(\d+)\.(?:key|value))"); + return re; +} + +// Matches e.g. "...layers.5.self_attn..." -> layer index 5. Reused from the same +// convention as embedding/prepare_embedding_model.cpp's AddKVCacheNodes matcher. +const std::regex& layer_id_regex() { + static const std::regex re(R"(layers\.(\d+)\.self_attn)"); + return re; +} + +bool try_parse_layer_idx(const std::string& text, const std::regex& re, size_t& out_idx) { + std::smatch m; + if (!std::regex_search(text, m, re)) { + return false; + } + out_idx = static_cast(std::stoul(m[1].str())); + return true; +} + +} // namespace + +namespace ov::npuw { + +PatchSlidingWindowKVCache::PatchSlidingWindowKVCache(uint32_t window_size, + std::vector layer_is_sliding, + uint32_t kvcache_size, + uint32_t input_size, + const KVAxesPosition& kv_axes_position, + bool trim_attention_mask) + : m_window_size(window_size), + m_layer_is_sliding(std::move(layer_is_sliding)), + m_kvcache_size(kvcache_size), + m_input_size(input_size), + m_kv_axes_position(kv_axes_position), + m_trim_attention_mask(trim_attention_mask) {} + +bool PatchSlidingWindowKVCache::run_on_model(const std::shared_ptr& model) { + if (m_window_size == 0 || m_layer_is_sliding.empty()) { + LOG_INFO("[SWA] Sliding Window Attention is not configured, skipping " << model->get_friendly_name()); + return false; + } + + // `new_kv_total` is the resulting post-concat K/V length seen by a sliding-window + // layer's SDPA (past + current input). The sliding-window causal constraint is + // per-query-row: the EARLIEST new query row in this call (absolute position P, + // i.e. tokens already stored before this call) needs past tokens + // [P-window_size+1, P-1], i.e. `window_size - 1` of history - independent of how + // many additional new rows (`m_input_size`) follow it in the same call. So past + // capacity must track `window_size` directly, NOT `window_size - input_size` + // (subtracting input_size here was a correctness bug: with input_size == window_size, + // e.g. chunked prefill, it collapsed past to 0, silently dropping all history the + // first row of the chunk should have attended to). + // Degenerate case: if window_size >= kvcache_size, sliding effectively never kicks + // in within the model's overall budget, so behave like a regular (non-SWA) layer: + // past = kvcache_size - input_size. + const uint32_t new_past_u32 = (m_window_size < m_kvcache_size) + ? m_window_size + : (m_input_size < m_kvcache_size ? (m_kvcache_size - m_input_size) : 0u); + const int64_t new_past = static_cast(new_past_u32); + const int64_t new_kv_total = new_past + static_cast(m_input_size); + bool changed = false; + + // Step 0: BEFORE touching any Parameter shape, snapshot the *current* (fully static, still + // uniform-`m_kvcache_size`) value of EVERY self_attn layer's shape/bound-carrying input that + // depends on the current KV length - both the GQA/repeat_kv `expand` lowering's target shape + // (Concat(past,new) -> Unsqueeze -> Broadcast -> Reshape -> SDPA) AND the per-layer causal-mask + // trim (HF's `causal_mask[..., : key_states.shape[-2]]`, lowered to `aten::slice` -> + // `opset8::Slice(mask, begin, end, step, axis)`, whose `end` bound is exactly this same + // current-KV-length value) - sliding *and* full-attention layers alike. This must happen + // strictly first, and must cover ALL layers, for two reasons: + // + // 1) `model->reshape()`/`validate_nodes_and_infer_types()` set-and-validate atomically, so + // any stale downstream shape needs to already be fixed *before* the single validation + // call at the end of this function - there's no chance to patch things "in between". + // + // 2) Some exporters/compilers common-subexpression-eliminate (CSE) the shape-computation + // subgraph feeding this input across MULTIPLE layers that happen to be structurally + // identical at trace time (all layers, sliding or full-attention, have the SAME uniform + // kvcache_size before this pass runs). Confirmed on a real Gemma model dump: one + // ShapeOf(Concat(past,new))->Gather->Concat chain, whose ShapeOf happened to be sourced + // from one PARTICULAR layer's own Concat, fed the Broadcast of SEVERAL other layers too + // (a mix of sliding AND full-attention ones). Only patching the sliding layers' consumer + // edges (and leaving full-attention layers alone, assuming they "don't need to change") + // is still wrong: once the representative source layer's Parameter is resized in Step 1, + // EVERY OTHER layer still wired to that shared node - including full-attention layers we + // never intended to touch - inherits the new (wrong-for-them) value too. This is exactly + // what caused a real crash where a *full-attention* layer's Broadcast target shape + // unexpectedly showed the *sliding* size, and then (same root cause) a *full-attention* + // layer's own causal-mask Slice's `end` bound too. So every layer's shape/bound input + // touching this pattern must be decoupled from the (possibly shared) upstream + // computation, regardless of that layer's own sliding status. + // + // The robust fix is to depend on neither of the above: right now, *before* any Parameter is + // reshaped, the whole model is still uniformly static, so `ov::util::get_constant_from_source()` + // can fold ANY such input - a plain Constant, or a live (possibly shared) ShapeOf-based chain - + // down to its current concrete value. We snapshot that value now, then later (Step 2, after + // Parameter shapes have changed) build a brand-new, per-node *private* Constant - for sliding + // layers with the stale `m_kvcache_size` entries corrected to `new_kv_total`, for + // full-attention layers with the SAME original value, just privatized - and hook it up via + // `replace_source_output()`. This never mutates or reuses the original (possibly shared) + // upstream node, so it works uniformly regardless of how the shape is represented or how many + // layers currently reuse the same shape-computation subgraph. + struct ShapeInputSnapshot { + std::shared_ptr consumer; + size_t input_index; + size_t layer_idx; + bool layer_is_sliding; + std::vector original_values; + }; + std::vector shape_snapshots; + // Nodes that ARE a per-layer causal-mask Slice (`causal_mask[..., begin:end]`) whose bound was + // found and will be privatized/corrected below. Tracked by NODE POINTER, not by the layer_idx + // parsed from the node's own friendly name: this Slice can ITSELF be CSE-shared across + // multiple layers' SDPA mask inputs (confirmed via a real crash - see Step 3's comment), so a + // single node whose name says "layer 0" may in fact be the actual mask source feeding many + // OTHER layers' SDPA too. For any SDPA whose mask input traces back to one of these nodes, the + // mask is ALREADY made consistent with the new KV length by this very mechanism - Step 3 + // further down must NOT also trim it (that would double-trim it). + std::unordered_set privatized_mask_slice_nodes; + for (const auto& node : model->get_ordered_ops()) { + // Reshape/Broadcast: target-shape input is input(1). Slice (the per-layer causal-mask + // trim `causal_mask[..., begin:end]`): both `begin` (input 1) and `end` (input 2) are + // guarded defensively, though in practice only `end` carries the current-KV-length value. + std::vector candidate_input_indices; + const bool is_mask_slice = ov::is_type(node); + if (ov::is_type(node) || ov::is_type(node)) { + candidate_input_indices = {1}; + } else if (is_mask_slice) { + candidate_input_indices = {1, 2}; + } else { + continue; + } + size_t layer_idx = 0; + if (!try_parse_layer_idx(node->get_friendly_name(), layer_id_regex(), layer_idx)) { + continue; + } + if (layer_idx >= m_layer_is_sliding.size()) { + continue; + } + for (const size_t input_idx : candidate_input_indices) { + if (input_idx >= node->get_input_size()) { + continue; + } + auto folded = ov::util::get_constant_from_source(node->input_value(input_idx)); + if (!folded) { + LOG_WARN("[SWA] Layer " << layer_idx << ": " << node->get_type_name() << " '" + << node->get_friendly_name() << "' input(" << input_idx + << ") is not constant-foldable, cannot verify/patch it."); + continue; + } + auto values = folded->cast_vector(); + const auto kv = static_cast(m_kvcache_size); + const bool has_kvcache_dim = + std::any_of(values.begin(), values.end(), [kv](int64_t v) { return v == kv || v == -kv; }); + if (!has_kvcache_dim) { + // Not a KV-size-dependent shape/bound (e.g. unrelated input in the same + // self_attn block) - nothing to guard here. + continue; + } + shape_snapshots.push_back( + ShapeInputSnapshot{node, input_idx, layer_idx, m_layer_is_sliding[layer_idx], std::move(values)}); + if (is_mask_slice) { + privatized_mask_slice_nodes.insert(node.get()); + } + } + } + + // Step 1: shrink past_key_values Parameter shapes for sliding-window layers. + // + // NOTE: we deliberately do NOT call `model->reshape()` here - see the rationale in Step 0's + // comment above. Instead we set the new Parameter shape directly + // (`Parameter::set_partial_shape`, which does not trigger any validation), and defer the + // single `validate_nodes_and_infer_types()` call to the very end of this function, once Step 2 + // has also fixed up the Broadcast/Reshape/Slice shape-bound inputs snapshotted in Step 0. + size_t num_params_reshaped = 0; + for (const auto& input : model->inputs()) { + const auto& name = input.get_any_name(); + size_t layer_idx = 0; + if (!try_parse_layer_idx(name, kv_param_regex(), layer_idx)) { + continue; + } + if (layer_idx >= m_layer_is_sliding.size() || !m_layer_is_sliding[layer_idx]) { + continue; + } + const auto& pshape = input.get_partial_shape(); + if (pshape.rank().is_dynamic() || m_kv_axes_position.seq_len >= pshape.size() || + !pshape[m_kv_axes_position.seq_len].is_static()) { + LOG_WARN("[SWA] Layer " << layer_idx << ": past KV parameter '" << name + << "' has a non-static seq_len axis, skipping reshape. Was ReshapeToStatic " + "applied before PatchSlidingWindowKVCache?"); + continue; + } + const int64_t old_past = pshape[m_kv_axes_position.seq_len].get_length(); + if (new_past == old_past) { + continue; + } + auto param = std::dynamic_pointer_cast(input.get_node_shared_ptr()); + if (!param) { + LOG_WARN("[SWA] Layer " << layer_idx << ": past KV input '" << name + << "' is not backed by a Parameter node, skipping."); + continue; + } + ov::PartialShape new_shape = pshape; + new_shape[m_kv_axes_position.seq_len] = new_past; + param->set_partial_shape(new_shape); + changed = true; + ++num_params_reshaped; + LOG_INFO("[SWA] Layer " << layer_idx << ": past KV '" << name << "' seq_len " << old_past << " -> " + << new_past << " (window=" << m_window_size << ", kvcache_size=" << m_kvcache_size + << ", input_size=" << m_input_size << ", post-concat total=" << new_kv_total + << ")"); + } + if (num_params_reshaped > 0) { + LOG_INFO("[SWA] Reshaped " << num_params_reshaped << " past_key_values parameter(s) in '" + << model->get_friendly_name() << "' for sliding-window layers."); + } else { + LOG_INFO("[SWA] No past_key_values parameters required a reshape in '" << model->get_friendly_name() + << "'."); + } + + // Step 2: apply the fix for every shape/bound input snapshotted in Step 0 (Broadcast/Reshape + // target shapes, and Slice begin/end bounds for the per-layer causal-mask trim). For each, + // take its ORIGINAL (pre-patch) value and rebuild it as a brand-new, PRIVATE Constant via + // `replace_source_output()` - never mutating the original (possibly shared) upstream node in + // place: + // - sliding-window layers: entries equal to `m_kvcache_size` (or `-m_kvcache_size`, for + // Slice bounds expressed as a negative/from-the-end offset) are corrected to + // `new_kv_total` (signed to match). + // - full-attention layers: entries are kept at their original value - but STILL rebuilt as a + // private Constant, purely to sever any (possibly shared) live upstream dependency that a + // sibling sliding layer might later invalidate (see Step 0's comment for why this matters + // even when the value itself doesn't change). + { + // Step 2a: for sliding-window layers' mask-Slice nodes specifically, the naive "correct the + // `end` bound constant" approach above is not enough. HF's `causal_mask[..., :new_kv_total]` + // lowers to a Slice with begin=0 (front-aligned), but the SWA runtime KV-write logic keeps the + // physical past_key_values buffer LEFT-aligned, holding the most recent tokens in chronological + // order (see the "KV buffer invariant" - oldest kept token at column 0, newest at the last + // column). So after this call, the buffer holds exactly ONE contiguous window of ABSOLUTE + // positions: `[chunk_start + input_size - new_kv_total, chunk_start + input_size)`, where + // `chunk_start` is this call's first query position (varies across calls that reuse the SAME + // compiled graph: different chunks of the same chunked-prefill model, or different steps served + // by the same generate kv_size variant). A single Slice with a FIXED begin (whether 0, as HF + // emits, or a fixed tail offset) can only ever be correct for ONE such call. + // + // Fix: replace the whole mask-Slice node with a Gather whose index tensor is + // `Range(0, new_kv_total, 1) + begin`, `begin` computed at runtime from `position_ids[0]`. + // NOTE: an earlier version of this fix used `Slice(begin, end)` with runtime bounds, pinned + // to a static shape via a `Reshape(special_zero=true)`. That crashed NPUW's online + // partitioning ("identifyUniques" pass -> "to_shape was called on a dynamic shape"): the + // model's own Parameter shapes are updated in Step 1 but downstream shape propagation is + // deferred to this function's own final `validate_nodes_and_infer_types()` call, so at the + // time this pass runs some non-sliced axes of the mask tensor are still legitimately dynamic + // (to be resolved by a LATER pipeline stage) - `special_zero` just copies that dynamism + // through, and the *sliced* axis's static pin doesn't help those OTHER axes. Gather doesn't + // have this problem: its output shape at the gathered axis comes purely from the INDEX + // tensor's OWN (always-static, `new_kv_total`-long) shape, and every other axis is passed + // through exactly as-is (dynamic or not) with no separate "pin" step required. + std::shared_ptr position_ids_node; + for (const auto& input : model->inputs()) { + if (input.get_any_name() == "position_ids") { + position_ids_node = input.get_node_shared_ptr(); + break; + } + } + const auto pos_ids_pshape = position_ids_node ? position_ids_node->get_output_partial_shape(0) : ov::PartialShape{}; + const bool position_ids_usable = + position_ids_node && pos_ids_pshape.rank().is_static() && pos_ids_pshape.size() == 2; + + // Lazily built, shared across layers. Kept as an explicit [1]-shaped (not fully-squeezed to a + // rank-0 scalar) tensor: the NPU/vpux compiler's own type inference for a Squeeze reducing this + // to a true scalar disagrees with its declared result type - Subtract/Maximum/Add below all + // broadcast a [1]-shaped operand just as well, so there's no need to squeeze it. + std::shared_ptr chunk_start_scalar; + auto get_chunk_start = [&]() { + if (!chunk_start_scalar) { + auto idx0 = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{}, {0}); + auto axis1 = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{1}, {1}); + chunk_start_scalar = std::make_shared(position_ids_node, idx0, axis1); + } + return chunk_start_scalar; + }; + + std::unordered_set dynamically_reselected; + size_t reselected_count = 0; + if (position_ids_usable && m_window_size < m_kvcache_size) { + for (const auto& snapshot : shape_snapshots) { + if (!snapshot.layer_is_sliding || !privatized_mask_slice_nodes.count(snapshot.consumer.get()) || + dynamically_reselected.count(snapshot.consumer.get())) { + continue; + } + auto slice_node = snapshot.consumer; + auto axis_const = ov::util::get_constant_from_source(slice_node->input_value(4)); + if (!axis_const) { + LOG_WARN("[SWA] Layer " << snapshot.layer_idx << ": mask Slice '" + << slice_node->get_friendly_name() + << "' has a non-constant-foldable axis, cannot dynamically reselect its " + "columns - falling back to the static-bound correction below."); + continue; + } + const int64_t axis = axis_const->cast_vector().at(0); + + const auto data_pshape = slice_node->get_input_partial_shape(0); + if (data_pshape.rank().is_dynamic()) { + LOG_WARN("[SWA] Layer " << snapshot.layer_idx << ": mask Slice '" + << slice_node->get_friendly_name() + << "' has a dynamic-rank data input, cannot dynamically reselect its " + "columns - falling back to the static-bound correction below."); + continue; + } + const int64_t rank = data_pshape.rank().get_length(); + const int64_t norm_axis = axis < 0 ? axis + rank : axis; + if (norm_axis < 0 || norm_axis >= rank) { + LOG_WARN("[SWA] Layer " << snapshot.layer_idx << ": mask Slice '" + << slice_node->get_friendly_name() << "' has an out-of-range axis " + << axis << " for rank " << rank + << " - falling back to the static-bound correction below."); + continue; + } + dynamically_reselected.insert(slice_node.get()); + + // begin_raw = chunk_start - window_size (UNCLAMPED, may be negative - this is the + // mathematically "true" start of the physical "past" region's absolute-position + // window). begin_1d = max(0, begin_raw) is only used to keep the Gather's own index + // values non-negative/valid; the possible discrepancy it introduces (see below) is + // compensated for explicitly, instead of silently accepted as before. + auto window_size_const = + ov::op::v0::Constant::create(ov::element::i64, ov::Shape{}, {static_cast(m_window_size)}); + auto begin_raw = std::make_shared(get_chunk_start(), window_size_const); + auto zero_scalar = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{}, {0}); + auto begin_1d = std::make_shared(begin_raw, zero_scalar); // shape [1] + + // The physical buffer is [past (window_size wide), new (input_size wide)] - reselect + // each part with its OWN index expression instead of one combined + // `Range(0, new_kv_total) + begin`: + // + // - "new" part: always `Range(0, input_size) + chunk_start`, UNCLAMPED. `chunk_start` + // (from position_ids) is never negative, so this is always valid, and it always + // picks exactly this call's own query positions - the columns holding this call's + // real, just-written K/V. + // + // - "past" part: `Range(0, window_size) + begin_1d`, clamped as before (Gather can't + // take negative indices). When `chunk_start >= window_size` (the steady-state case) + // `begin_raw` is already >= 0, so this is identical to the previous single-Gather + // behavior. But when `chunk_start < window_size` (only possible for the very FIRST + // call using this compiled graph, e.g. the first chunk of chunked prefill, since + // position_ids only grows afterwards), `begin_raw` is negative and clamping to 0 + // shifts the "past" part's gathered columns to alias the SAME absolute positions + // the "new" part just picked - silently making not-yet-existing "past" slots look + // like valid, visible duplicates of the current chunk's own tokens (instead of the + // intended "no real history yet" - invisible). Since the physical past_key_values + // buffer content at those slots is never actually written for such a call, this + // physically-nonexistent-but-visible aliasing let attention draw on garbage/stale + // KV data - the discrepancy this whole block corrects. + // + // Fix: explicitly force any "past" slot whose UNCLAMPED absolute position + // (`begin_raw + local_p`) is still negative back to invisible (-inf), regardless of + // which (borrowed/aliased) column the clamped Gather happened to read. + auto idx_new_start = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{}, {0}); + auto idx_new_stop = + ov::op::v0::Constant::create(ov::element::i64, ov::Shape{}, {static_cast(m_input_size)}); + auto idx_step = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{}, {1}); + auto idx_new_range = + std::make_shared(idx_new_start, idx_new_stop, idx_step, ov::element::i64); + auto new_indices = + std::make_shared(idx_new_range, get_chunk_start()); // shape [input_size] + + auto idx_past_start = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{}, {0}); + auto idx_past_stop = + ov::op::v0::Constant::create(ov::element::i64, ov::Shape{}, {static_cast(m_window_size)}); + auto idx_past_range = + std::make_shared(idx_past_start, idx_past_stop, idx_step, ov::element::i64); + auto past_indices = std::make_shared(idx_past_range, begin_1d); // shape [window_size] + // Same per-slot offsets, but against the UNCLAMPED begin - negative entries mark + // "past" slots that don't correspond to any real history yet. + auto past_raw_pos = std::make_shared(idx_past_range, begin_raw); // shape [window_size] + auto invalid_past = + std::make_shared(past_raw_pos, zero_scalar); // shape [window_size], BOOL + + auto axis_1d = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{1}, {axis}); + auto new_gathered = + std::make_shared(slice_node->input_value(0), new_indices, axis_1d); + auto past_gathered = + std::make_shared(slice_node->input_value(0), past_indices, axis_1d); + + // Broadcast `invalid_past` ([window_size]) to the mask's rank, with `window_size` at + // `norm_axis` and `1` everywhere else, so it can Select elementwise against + // `past_gathered` (whose other axes may still be dynamic at this point - the reshape + // target itself is fully static, computed purely from `rank`/`norm_axis`/`window_size`). + std::vector invalid_shape(static_cast(rank), 1); + invalid_shape[static_cast(norm_axis)] = static_cast(m_window_size); + auto invalid_shape_const = ov::op::v0::Constant::create(ov::element::i64, + ov::Shape{invalid_shape.size()}, + invalid_shape); + auto invalid_past_reshaped = + std::make_shared(invalid_past, invalid_shape_const, false); + + const auto data_elem_type = slice_node->get_input_element_type(0); + auto neg_inf_const = + ov::op::v0::Constant::create(data_elem_type, ov::Shape{}, {-std::numeric_limits::max()}); + auto past_corrected = + std::make_shared(invalid_past_reshaped, neg_inf_const, past_gathered); + + auto gathered = std::make_shared( + ov::OutputVector{past_corrected, new_gathered}, + axis); + gathered->set_friendly_name(slice_node->get_friendly_name() + "/swa_dynamic_reselect"); + + // Keep Step 3's "already handled" guard working: it looks up an SDPA's mask source + // node in `privatized_mask_slice_nodes` to skip double-trimming. Since consumers now + // point at `gathered` instead of the original Slice, register it too. + privatized_mask_slice_nodes.insert(gathered.get()); + + auto target_inputs = slice_node->output(0).get_target_inputs(); + for (auto&& input : target_inputs) { + input.replace_source_output(gathered); + } + changed = true; + ++reselected_count; + LOG_INFO("[SWA] Layer " << snapshot.layer_idx << ": replaced mask Slice '" + << slice_node->get_friendly_name() + << "' with a position_ids-anchored dynamic-begin Gather (window=" + << m_window_size << ", input_size=" << m_input_size + << ", total_width=" << m_kvcache_size << ")."); + } + } + if (reselected_count > 0) { + LOG_INFO("[SWA] Dynamically reselected " << reselected_count + << " sliding-layer mask Slice node(s) in '" + << model->get_friendly_name() << "' using position_ids."); + } + + // Step 2b: the original constant-correction path, for everything NOT handled by Step 2a above + // (Broadcast/Reshape shape-bound inputs, full-attention-layer mask-Slice bounds, and + // sliding-layer mask-Slice bounds in the degenerate `window_size >= kvcache_size` case or when + // `position_ids` wasn't usable). + size_t patched_count = 0; + size_t privatized_count = 0; + for (const auto& snapshot : shape_snapshots) { + if (dynamically_reselected.count(snapshot.consumer.get())) { + continue; // node fully replaced above; its old begin/end constants are now dead code. + } + auto values = snapshot.original_values; + const int64_t kv = static_cast(m_kvcache_size); + const int64_t corrected = snapshot.layer_is_sliding ? new_kv_total : kv; + bool value_changed = false; + for (auto& v : values) { + if (v == kv && v != corrected) { + v = corrected; + value_changed = true; + } else if (v == -kv && v != -corrected) { + v = -corrected; + value_changed = true; + } + } + auto new_const = std::make_shared(ov::element::i64, ov::Shape{values.size()}, values); + new_const->set_friendly_name(snapshot.consumer->get_friendly_name() + "/swa_shape_patched_" + + std::to_string(snapshot.input_index)); + snapshot.consumer->input(snapshot.input_index).replace_source_output(new_const); + changed = true; + ++privatized_count; + if (value_changed) { + ++patched_count; + LOG_INFO("[SWA] Layer " << snapshot.layer_idx << ": patched " << snapshot.consumer->get_type_name() + << " '" << snapshot.consumer->get_friendly_name() << "' input(" + << snapshot.input_index << ") constant " << m_kvcache_size << " -> " + << new_kv_total); + } + } + if (privatized_count > 0) { + LOG_INFO("[SWA] Privatized " << privatized_count + << " Broadcast/Reshape/Slice shape-bound input(s) (" << patched_count + << " value-corrected, " << (privatized_count - patched_count) + << " guarded-only) feeding self_attn SDPA node(s) in '" + << model->get_friendly_name() << "'."); + } + } + + // Step 3 (generate model only): trim the SDPA attention-mask input to the same window, for + // architectures where the mask reaching SDPA is NOT already per-layer trimmed to the current + // KV length by the model itself. + // + // IMPORTANT: this must be SKIPPED for any SDPA whose mask input traces back to a node in + // `privatized_mask_slice_nodes` (Step 0/2's causal-mask Slice privatization). Confirmed via a + // real crash: for architectures that DO lower `causal_mask[..., :key_states.shape[-2]]` to a + // per-layer `opset8::Slice`, Step 0/2 already corrects that Slice's own `end` bound to + // `new_kv_total`, so the SDPA mask input is already the right width. Step 3 running on TOP of + // that (still) used a STALE `mask_source.get_partial_shape()` - queried BEFORE the deferred + // final `validate_nodes_and_infer_types()` - so it saw the OLD (pre-Step-2) width and inserted + // an ADDITIONAL Slice on top, double-trimming the mask down to a wrong, too-narrow width + // (observed: 898 instead of the correct new_kv_total). + // + // A SECOND, MORE SUBTLE layer to this same bug (also confirmed via a real crash): the mask + // Slice node itself can be CSE-shared across MULTIPLE layers (exactly like the Broadcast/ + // Reshape sharing described in Step 0) - e.g. only ONE Slice node, whose friendly name happens + // to say "layer 0", is in fact the actual mask source for 24 OTHER sliding layers' SDPA too. + // So the guard here CANNOT be keyed by "the layer_idx parsed from the Slice node's own name" + // (that only matches the one layer whose name the shared node happens to carry) - it must be + // keyed by the ACTUAL node identity of `mask_source.get_node()`, checked against every node we + // privatized in Step 0, regardless of which layer's name that node carries. + if (m_trim_attention_mask) { + std::unordered_map> mask_slice_cache; + size_t sdpa_count = 0; + size_t patched_count = 0; + for (const auto& node : model->get_ordered_ops()) { + auto sdpa = std::dynamic_pointer_cast(node); + if (!sdpa) { + continue; + } + ++sdpa_count; + size_t layer_idx = 0; + if (!try_parse_layer_idx(sdpa->get_friendly_name(), layer_id_regex(), layer_idx)) { + LOG_INFO("[SWA] SDPA node '" << sdpa->get_friendly_name() + << "' has no parsable layer index, skipping mask trim."); + continue; + } + if (layer_idx >= m_layer_is_sliding.size() || !m_layer_is_sliding[layer_idx]) { + continue; + } + static constexpr size_t kMaskInputIdx = 3; // Q=0, K=1, V=2, mask=3 (see attention.cpp SDPA_Inputs) + if (sdpa->get_input_size() <= kMaskInputIdx) { + LOG_INFO("[SWA] SDPA node '" << sdpa->get_friendly_name() << "' (layer " << layer_idx + << ") has no attention-mask input, skipping mask trim."); + continue; + } + auto mask_input = sdpa->input(kMaskInputIdx); + const auto mask_source = mask_input.get_source_output(); + if (privatized_mask_slice_nodes.count(mask_source.get_node()) > 0) { + LOG_INFO("[SWA] SDPA node '" + << sdpa->get_friendly_name() << "' (layer " << layer_idx << ") mask source '" + << mask_source.get_node()->get_friendly_name() + << "' is already trimmed to the current KV length via Step 2's causal-mask Slice " + "privatization - skipping redundant mask trim."); + continue; + } + const auto mask_pshape = mask_source.get_partial_shape(); + if (mask_pshape.rank().is_dynamic() || mask_pshape.size() == 0 || + !mask_pshape[mask_pshape.size() - 1].is_static()) { + LOG_WARN("[SWA] SDPA node '" << sdpa->get_friendly_name() << "' (layer " << layer_idx + << ") has a dynamic mask shape, cannot trim."); + continue; + } + const size_t last_axis = mask_pshape.size() - 1; + const int64_t old_width = mask_pshape[last_axis].get_length(); + if (new_kv_total >= old_width) { + // Mask is already narrow enough (e.g. window >= kvcache_size for this variant). + continue; + } + + ov::Output sliced; + auto cache_it = mask_slice_cache.find(mask_source.get_node()); + if (cache_it != mask_slice_cache.end()) { + sliced = cache_it->second; + } else { + auto begin = ov::op::v0::Constant::create(ov::element::i64, {1}, {old_width - new_kv_total}); + auto end = ov::op::v0::Constant::create(ov::element::i64, {1}, {old_width}); + auto step = ov::op::v0::Constant::create(ov::element::i64, {1}, {1}); + auto axis = ov::op::v0::Constant::create(ov::element::i64, {1}, {static_cast(last_axis)}); + auto slice = std::make_shared(mask_source, begin, end, step, axis); + slice->set_friendly_name(mask_source.get_node()->get_friendly_name() + "/swa_mask_slice"); + sliced = slice->output(0); + mask_slice_cache.emplace(mask_source.get_node(), sliced); + LOG_INFO("[SWA] Inserted shared mask Slice in '" + << model->get_friendly_name() << "': width " << old_width << " -> " << new_kv_total + << " (source node: " << mask_source.get_node()->get_friendly_name() << ")"); + } + mask_input.replace_source_output(sliced); + ++patched_count; + changed = true; + } + LOG_INFO("[SWA] '" << model->get_friendly_name() << "': scanned " << sdpa_count + << " SDPA node(s), patched mask input on " << patched_count + << " sliding-layer SDPA node(s) using " << mask_slice_cache.size() + << " unique Slice node(s)."); + } + + if (changed) { + model->validate_nodes_and_infer_types(); + } + return changed; +} + +} // namespace ov::npuw diff --git a/src/plugins/intel_npu/src/plugin/npuw/npuw_transformations/patch_sliding_window_kvcache.hpp b/src/plugins/intel_npu/src/plugin/npuw/npuw_transformations/patch_sliding_window_kvcache.hpp new file mode 100644 index 00000000000000..55ef1465159c81 --- /dev/null +++ b/src/plugins/intel_npu/src/plugin/npuw/npuw_transformations/patch_sliding_window_kvcache.hpp @@ -0,0 +1,53 @@ +// Copyright (C) 2018-2026 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#pragma once + +#include +#include + +#include "kv_axes_position.hpp" +#include "openvino/pass/pass.hpp" + +namespace ov::npuw { + +// Reduces the past_key_values buffer size (and, optionally, the SDPA attention mask +// width) for layers configured as Sliding Window Attention (SWA), e.g. Gemma4's +// "sliding_attention" layers. Must run AFTER ReshapeToStatic (it re-shrinks a subset +// of already-static past_key_values Parameters) and BEFORE DecomposeGQA / value-tensor +// optimization passes. +// +// Step 1 (always performed): for every past_key_values..key/value Parameter that +// belongs to a sliding-window layer, shrink its seq_len axis from +// (kvcache_size - input_size) down to (min(window_size, kvcache_size) - input_size), +// clamped to 0. This is what actually reduces per-layer NPU memory/compute for SWA. +// +// Step 2 (only when trim_attention_mask == true): for the same sliding-window layers, +// insert a shared `Slice` on the SDPA node's attention-mask input (port 3), trimming its +// last (key) axis down to the same `min(window_size, kvcache_size)` width. This is only +// valid for models where every query position shares the same trailing K/V window - i.e. +// the *generate* model (decode step, `max_generation_token_len` new tokens). It must NOT +// be applied to the *prefill* model, where multiple query positions attend to different, +// staggered windows over the same K/V span (that per-token banded pattern is produced by +// the existing SlidingWindowMask pass and must not be blanket-trimmed here). +class PatchSlidingWindowKVCache : public ov::pass::ModelPass { + uint32_t m_window_size; + std::vector m_layer_is_sliding; + uint32_t m_kvcache_size; + uint32_t m_input_size; + KVAxesPosition m_kv_axes_position; + bool m_trim_attention_mask; + +public: + OPENVINO_MODEL_PASS_RTTI("ov::npuw::PatchSlidingWindowKVCache"); + PatchSlidingWindowKVCache(uint32_t window_size, + std::vector layer_is_sliding, + uint32_t kvcache_size, + uint32_t input_size, + const KVAxesPosition& kv_axes_position, + bool trim_attention_mask); + bool run_on_model(const std::shared_ptr& model) override; +}; + +} // namespace ov::npuw diff --git a/src/plugins/intel_npu/src/plugin/npuw/npuw_transformations/sliding_window_mask.cpp b/src/plugins/intel_npu/src/plugin/npuw/npuw_transformations/sliding_window_mask.cpp index 510d1076382819..4af78efbbe4a86 100644 --- a/src/plugins/intel_npu/src/plugin/npuw/npuw_transformations/sliding_window_mask.cpp +++ b/src/plugins/intel_npu/src/plugin/npuw/npuw_transformations/sliding_window_mask.cpp @@ -645,6 +645,80 @@ class Gemma4UnifiedSlidingMaskMatcher : public ov::pass::MatcherPass { } }; +class Gemma4MaskedFillSlidingMaskMatcher : public ov::pass::MatcherPass { +public: + OPENVINO_MATCHER_PASS_RTTI("ov::npuw::patterns::Gemma4MaskedFillSlidingMaskMatcher"); + + explicit Gemma4MaskedFillSlidingMaskMatcher(const std::shared_ptr& position_ids_node_ptr) { + // Fixes a THIRD, newer Gemma-4 sliding-window mask export shape (confirmed on a real + // chunked-prefill dump), built via "masked_fill" (lowered to Select) instead of the + // BitwiseAnd-based forms handled by Gemma4SlidingMaskMatcher / Gemma4UnifiedSlidingMaskMatcher + // above: + // + // col_pos = Unsqueeze(Range(0, kvcache_size, 1)) -- [1, kvcache_size], static + // row_pos = Add(Unsqueeze(Range(0, chunk_size, 1)), past_len) -- [chunk_size, 1] + // beyond_window = GreaterEqual(Subtract(row_pos, col_pos), window_size) + // sliding_mask = Select(Unsqueeze(Unsqueeze(beyond_window)), -inf, causal_mask) + // + // `past_len` here is a plain compile-time Constant (observed to fold to + // `kvcache_size - chunk_size`, i.e. it silently assumes every call's queries sit at the LAST + // `chunk_size` absolute positions of the buffer) - it never reads `position_ids`. For a model + // reused across multiple chunks/generate steps (chunked prefill, or any generate KV-size + // variant), the window boundary is therefore frozen to whichever single call happens to match + // that assumption, silently corrupting `beyond_window`/`sliding_mask` for every other call. + // + // Fix: replace the `past_len` operand of `row_pos`'s Add with the REAL current chunk start, + // read from `position_ids[0]` at runtime. This only fixes mask CONTENT - the separate + // width-selection Slice (patched by PatchSlidingWindowKVCache) still needs its own, + // position_ids-anchored column-selection fix. + auto col_start_const = opp::wrap_type(); + auto key_range = opp::wrap_type({col_start_const, opp::any_input(), opp::any_input()}); + auto col_pos = opp::wrap_type({key_range, opp::any_input()}); + + auto row_start_const = opp::wrap_type(); + auto local_arange = opp::wrap_type({row_start_const, opp::any_input(), opp::any_input()}); + auto local_arange_row = opp::wrap_type({local_arange, opp::any_input()}); + auto row_pos = opp::wrap_type({local_arange_row, opp::any_input()}); + + auto row_minus_col = opp::wrap_type({row_pos, col_pos}); + auto beyond_window = opp::wrap_type({row_minus_col, opp::any_input()}); + auto beyond_window_u1 = opp::wrap_type({beyond_window, opp::any_input()}); + auto beyond_window_u2 = opp::wrap_type({beyond_window_u1, opp::any_input()}); + auto sliding_mask = opp::wrap_type({beyond_window_u2, opp::any_input(), opp::any_input()}); + + auto callback = [=](opp::Matcher& m) { + auto& node_to_output = m.get_pattern_value_map(); + auto row_pos_node = node_to_output.at(row_pos).get_node_shared_ptr(); + + // chunk_start = position_ids[:, 0:1] -- the real absolute position of the first query row + // in THIS call, read at runtime. Unlike PatchSlidingWindowKVCache's mask-column reselection + // (which needs a runtime-computed begin), the index here is a compile-time-literal 0, so a + // plain `Slice` already yields a fully static output shape - cheaper than a `Gather` and + // simpler (no Reshape-pin needed). Slice (unlike Gather with a scalar index) keeps the + // sliced axis instead of removing it, so the result stays [batch, 1] rather than [batch] - + // Add below broadcasts a [batch, 1]-shaped operand just as well. + auto begin_1d = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{1}, {0}); + auto end_1d = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{1}, {1}); + auto step_1d = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{1}, {1}); + auto axis_1d = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{1}, {1}); + std::shared_ptr chunk_start = + std::make_shared(position_ids_node_ptr, begin_1d, end_1d, step_1d, axis_1d); + + const auto target_et = row_pos_node->get_input_element_type(1); + if (chunk_start->get_element_type() != target_et) { + chunk_start = std::make_shared(chunk_start, target_et); + } + row_pos_node->input(1).replace_source_output(chunk_start); + + LOG_INFO("Found Gemma-4 (masked_fill-style) Sliding Window Attention mask, replaced its static " + "past-length reference with a position_ids-derived value."); + return true; + }; + register_matcher(std::make_shared(sliding_mask, "Gemma4MaskedFillSlidingMaskMatcher"), + std::move(callback)); + } +}; + #ifdef __GNUC__ # pragma GCC diagnostic pop #endif @@ -680,6 +754,7 @@ bool SlidingWindowMask::run_on_model(const std::shared_ptr& model) { const auto rewriter = manager.register_pass(); rewriter->add_matcher(attention_mask_node_ptr, position_ids_node_ptr); rewriter->add_matcher(attention_mask_node_ptr, position_ids_node_ptr); + rewriter->add_matcher(position_ids_node_ptr); rewriter->add_matcher(attention_mask_node_ptr, position_ids_node_ptr); rewriter->add_matcher(); return manager.run_passes(model); diff --git a/src/plugins/intel_npu/src/plugin/sources.cmake b/src/plugins/intel_npu/src/plugin/sources.cmake index b7a7c3e4fd6b02..25bef90b98584c 100644 --- a/src/plugins/intel_npu/src/plugin/sources.cmake +++ b/src/plugins/intel_npu/src/plugin/sources.cmake @@ -131,6 +131,8 @@ set(NPUW_SOURCES ${CMAKE_CURRENT_SOURCE_DIR}/npuw/npuw_transformations/lora_stateful_to_stateless.hpp ${CMAKE_CURRENT_SOURCE_DIR}/npuw/npuw_transformations/optimize_value_tensors.cpp ${CMAKE_CURRENT_SOURCE_DIR}/npuw/npuw_transformations/optimize_value_tensors.hpp + ${CMAKE_CURRENT_SOURCE_DIR}/npuw/npuw_transformations/patch_sliding_window_kvcache.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/npuw/npuw_transformations/patch_sliding_window_kvcache.hpp ${CMAKE_CURRENT_SOURCE_DIR}/npuw/npuw_transformations/patch_sliding_window_mask.cpp ${CMAKE_CURRENT_SOURCE_DIR}/npuw/npuw_transformations/patch_sliding_window_mask.hpp ${CMAKE_CURRENT_SOURCE_DIR}/npuw/npuw_transformations/remove_token_type_ids.cpp diff --git a/src/plugins/intel_npu/tests/unit/CMakeLists.txt b/src/plugins/intel_npu/tests/unit/CMakeLists.txt index f15da2723d61b3..aed04d936239a0 100644 --- a/src/plugins/intel_npu/tests/unit/CMakeLists.txt +++ b/src/plugins/intel_npu/tests/unit/CMakeLists.txt @@ -110,6 +110,7 @@ ov_add_test_target( ${OpenVINO_SOURCE_DIR}/src/plugins/intel_npu/src/plugin/npuw/npuw_transformations/untangle_dq_scale.cpp ${OpenVINO_SOURCE_DIR}/src/plugins/intel_npu/src/plugin/npuw/npuw_transformations/lora_stateful_to_stateless.cpp ${OpenVINO_SOURCE_DIR}/src/plugins/intel_npu/src/plugin/npuw/npuw_transformations/optimize_value_tensors.cpp + ${OpenVINO_SOURCE_DIR}/src/plugins/intel_npu/src/plugin/npuw/npuw_transformations/patch_sliding_window_kvcache.cpp ${OpenVINO_SOURCE_DIR}/src/plugins/intel_npu/src/plugin/npuw/npuw_transformations/patch_sliding_window_mask.cpp ${OpenVINO_SOURCE_DIR}/src/plugins/intel_npu/src/plugin/npuw/npuw_transformations/sliding_window_mask.cpp ${OpenVINO_SOURCE_DIR}/src/plugins/intel_npu/src/plugin/npuw/npuw_transformations/detect_causal_mask.cpp @@ -160,6 +161,7 @@ target_sources(${TARGET_NAME} PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/npuw/llm_continuation_test.cpp ${CMAKE_CURRENT_SOURCE_DIR}/npuw/stored_tokens_state_test.cpp ${CMAKE_CURRENT_SOURCE_DIR}/npuw/llm_infer_request_variant_switch_test.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/npuw/write_kv_slice_sliding_test.cpp ${CMAKE_CURRENT_SOURCE_DIR}/npuw/llm_trim_kvcache_test.cpp ${CMAKE_CURRENT_SOURCE_DIR}/npuw/lincache_utils_test.cpp ${CMAKE_CURRENT_SOURCE_DIR}/npuw/model_builder_lora_test.cpp diff --git a/src/plugins/intel_npu/tests/unit/npuw/host_flash_attention_test.cpp b/src/plugins/intel_npu/tests/unit/npuw/host_flash_attention_test.cpp index f005b12fe7f5d3..dba92c2eeb142e 100644 --- a/src/plugins/intel_npu/tests/unit/npuw/host_flash_attention_test.cpp +++ b/src/plugins/intel_npu/tests/unit/npuw/host_flash_attention_test.cpp @@ -9,6 +9,7 @@ #include #include +#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 add; + for (const auto& node : model->get_ops()) { + add = ov::as_type_ptr(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(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 add; + for (const auto& node : model->get_ops()) { + add = ov::as_type_ptr(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(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 // ============================================================================ diff --git a/src/plugins/intel_npu/tests/unit/npuw/pipeline_passes/detect_causal_mask_test.cpp b/src/plugins/intel_npu/tests/unit/npuw/pipeline_passes/detect_causal_mask_test.cpp index b1008e353742fa..dc32168cc7a50b 100644 --- a/src/plugins/intel_npu/tests/unit/npuw/pipeline_passes/detect_causal_mask_test.cpp +++ b/src/plugins/intel_npu/tests/unit/npuw/pipeline_passes/detect_causal_mask_test.cpp @@ -7,6 +7,7 @@ #include #include +#include #include "../llm_test_helpers.hpp" #include "model_builder.hpp" @@ -14,6 +15,7 @@ #include "openvino/op/scaled_dot_product_attention.hpp" #include "openvino/openvino.hpp" +using ov::npuw::AnnotatePerSDPAMaskType; using ov::npuw::DetectAttentionMask; using ov::npuw::MaskInfo; using MaskType = ov::npuw::MaskInfo::MaskType; @@ -32,6 +34,105 @@ MaskType detect(const std::shared_ptr& model) { return pass.get_mask_info().mask_type; } +std::shared_ptr append_decomposed_sdpa_branch(int layer_idx, + bool is_sliding_mask, + ov::ParameterVector& params, + ov::ResultVector& results) { + using namespace ov::op; + + const std::string idx = std::to_string(layer_idx); + + auto q = std::make_shared(ov::element::f32, ov::Shape{1, 1, 4, 8}); + auto past_k = std::make_shared(ov::element::f32, ov::Shape{1, 1, 4, 8}); + auto present_k = std::make_shared(ov::element::f32, ov::Shape{1, 1, 4, 8}); + auto past_v = std::make_shared(ov::element::f32, ov::Shape{1, 1, 4, 8}); + auto present_v = std::make_shared(ov::element::f32, ov::Shape{1, 1, 4, 8}); + + q->set_friendly_name("query." + idx); + past_k->set_friendly_name("past_key_values." + idx + ".key"); + present_k->set_friendly_name("present." + idx + ".key"); + past_v->set_friendly_name("past_key_values." + idx + ".value"); + present_v->set_friendly_name("present." + idx + ".value"); + + params.insert(params.end(), {q, past_k, present_k, past_v, present_v}); + + auto key_concat = std::make_shared(ov::OutputVector{past_k, present_k}, 2); + key_concat->set_friendly_name("concat_key." + idx); + auto value_concat = std::make_shared(ov::OutputVector{past_v, present_v}, 2); + value_concat->set_friendly_name("concat_value." + idx); + + auto zero_i64 = v0::Constant::create(ov::element::i64, ov::Shape{}, {0}); + auto one_i64 = v0::Constant::create(ov::element::i64, ov::Shape{}, {1}); + auto four_i64 = v0::Constant::create(ov::element::i64, ov::Shape{}, {4}); // Q seq length + auto eight_i64 = + v0::Constant::create(ov::element::i64, ov::Shape{}, {8}); // KV context length (past=4 + present=4) + + // k_range covers all KV positions; q_range covers query positions only. + // k_unsq → [1, kv_len], q_unsq → [seq, 1]; comparison broadcasts to [seq, kv_len]. + auto k_range = std::make_shared(zero_i64, eight_i64, one_i64, ov::element::i64); + auto q_range = std::make_shared(zero_i64, four_i64, one_i64, ov::element::i64); + auto k_unsq = std::make_shared(k_range, zero_i64); // [1, 8] + auto q_unsq = std::make_shared(q_range, one_i64); // [4, 1] + + std::shared_ptr mask_bool; + if (is_sliding_mask) { + auto neg_window = v0::Constant::create(ov::element::i64, ov::Shape{}, {-2}); + auto bound = std::make_shared(q_unsq, neg_window); + auto greater = std::make_shared(k_unsq, bound); + auto causal = std::make_shared(k_unsq, q_unsq); + auto one_bool = v0::Constant::create(ov::element::boolean, ov::Shape{}, {true}); + auto and_win = std::make_shared(one_bool, greater); + mask_bool = std::make_shared(and_win, causal); + } else { + mask_bool = std::make_shared(k_unsq, q_unsq); + } + + auto zero_f = v0::Constant::create(ov::element::f32, ov::Shape{}, {0.0f}); + auto neg_inf = v0::Constant::create(ov::element::f32, ov::Shape{}, {-std::numeric_limits::infinity()}); + auto mask_f = std::make_shared(mask_bool, zero_f, neg_inf); + auto m0 = std::make_shared(mask_f, zero_i64); + auto mask_4d = std::make_shared(m0, zero_i64); + + auto qk = std::make_shared(q, key_concat, false, true); + qk->set_friendly_name("matmul1." + idx); + auto add = std::make_shared(qk, mask_4d); + add->set_friendly_name("add." + idx); + auto softmax = std::make_shared(add, -1); + softmax->set_friendly_name("softmax." + idx); + auto out = std::make_shared(softmax, value_concat, false, false); + out->set_friendly_name("matmul2." + idx); + + auto result = std::make_shared(out); + result->set_friendly_name("attn_out." + idx); + results.push_back(result); + return add; +} + +std::pair, std::shared_ptr> make_decomposed_sdpa_with_mask(bool sliding_mask) { + ov::ParameterVector params; + ov::ResultVector results; + auto add = append_decomposed_sdpa_branch(/*layer_idx=*/0, sliding_mask, params, results); + auto model = std::make_shared(results, params); + return {model, add}; +} + +struct MixedAttentionModel { + std::shared_ptr model; + std::shared_ptr add_global; + std::shared_ptr add_swa; +}; + +MixedAttentionModel make_mixed_decomposed_sdpa_model() { + ov::ParameterVector params; + ov::ResultVector results; + + auto add_global = append_decomposed_sdpa_branch(/*layer_idx=*/0, /*is_sliding_mask=*/false, params, results); + auto add_swa = append_decomposed_sdpa_branch(/*layer_idx=*/1, /*is_sliding_mask=*/true, params, results); + + auto model = std::make_shared(results, params); + return {model, add_global, add_swa}; +} + // --------------------------------------------------------------------------- // Minimal Q/K/V SDPA model (no explicit mask). is_causal drives the op attribute. // --------------------------------------------------------------------------- @@ -177,6 +278,65 @@ std::shared_ptr build_minicpm_less() { return std::make_shared(ov::ResultVector{result}, ov::ParameterVector{seq, amask}); } +// --------------------------------------------------------------------------- +// Faithful reproduction of the THIRD, newer Gemma-4 "masked_fill" sliding +// window mask export shape (see Gemma4MaskedFillSlidingMaskMatcher in +// sliding_window_mask.cpp). The window check is a single GreaterEqual +// combined via Select/masked_fill, and — crucially — the causal part is a +// SEPARATE, standalone LessEqual(range_chain, range_chain) sub-comparison +// (the same shape StandardCausalMatcher recognizes on its own): +// +// col_pos = Unsqueeze(Range(0, total, 1)) +// row_pos = Add(Unsqueeze(Range(0, seq, 1)), past) +// beyond_window = GreaterEqual(Subtract(row_pos, col_pos), window) +// causal_mask = LessEqual(kv_row, q_col) -- independent chain +// sliding_mask = Select(Unsqueeze(Unsqueeze(beyond_window)), -inf, causal_mask) +// --------------------------------------------------------------------------- +std::shared_ptr build_gemma4_masked_fill_sliding(int64_t window) { + using namespace ov::op; + auto seq = std::make_shared(ov::element::i64, ov::PartialShape{1, -1}); + auto amask = std::make_shared(ov::element::i64, ov::PartialShape{1, -1}); + + auto zero = v0::Constant::create(ov::element::i64, ov::Shape{}, {0}); + auto one = v0::Constant::create(ov::element::i64, ov::Shape{}, {1}); + + auto ids_shape = std::make_shared(seq, ov::element::i64); + auto mask_shape = std::make_shared(amask, ov::element::i64); + auto seq_len = std::make_shared(ids_shape, one, zero); + auto total = std::make_shared(mask_shape, one, zero); + auto past = std::make_shared(total, seq_len); + + // col_pos: Range(0, total) -> Unsqueeze -> [1, total] + auto key_range = std::make_shared(zero, total, one, ov::element::i64); + auto col_pos = std::make_shared(key_range, zero); + + // row_pos: Add(Unsqueeze(Range(0, seq)), past) -> [seq, 1] + auto local_arange = std::make_shared(zero, seq_len, one, ov::element::i64); + auto local_arange_row = std::make_shared(local_arange, one); + auto row_pos = std::make_shared(local_arange_row, past); + + auto row_minus_col = std::make_shared(row_pos, col_pos); + auto window_const = v0::Constant::create(ov::element::i64, ov::Shape{}, {window}); + auto beyond_window = std::make_shared(row_minus_col, window_const); + auto bw_u1 = std::make_shared(beyond_window, zero); + auto bw_u2 = std::make_shared(bw_u1, zero); + + // Separate causal sub-comparison, own Range chain, matching StandardCausalMatcher. + auto kv_range = std::make_shared(zero, total, one, ov::element::i64); + auto kv_row = std::make_shared(kv_range, zero); + auto q_range = std::make_shared(zero, seq_len, one, ov::element::i64); + auto q_abs = std::make_shared(q_range, past); + auto q_col = std::make_shared(q_abs, one); + auto causal_mask = std::make_shared(kv_row, q_col); + + auto zero_f = v0::Constant::create(ov::element::f32, ov::Shape{}, {0.0f}); + auto neg_inf = v0::Constant::create(ov::element::f32, ov::Shape{}, {-std::numeric_limits::infinity()}); + auto causal_mask_f = std::make_shared(causal_mask, zero_f, neg_inf); + auto sliding_mask = std::make_shared(bw_u2, neg_inf, causal_mask_f); + auto result = std::make_shared(sliding_mask); + return std::make_shared(ov::ResultVector{result}, ov::ParameterVector{seq, amask}); +} + } // namespace // ============================================================================ @@ -378,12 +538,157 @@ TEST(DetectAttentionMaskTest, RealPattern_Phi3Sliding_IsSlidingWindowWithSize) { EXPECT_EQ(info.window_size, window); } +// Real pattern: newer Gemma-4 "masked_fill"-style SWA export (GreaterEqual + +// Select, with a standalone causal LessEqual sub-comparison). Regression test +// for the bug where the standalone causal sub-comparison alone got picked up +// by StandardCausalMatcher, mis-tagging the whole model as plain Causal and +// silently disabling the HFA explicit mask for SWA layers. +TEST(DetectAttentionMaskTest, RealPattern_Gemma4MaskedFillSliding_IsSlidingWindowWithSize) { + const int64_t window = 1024; + DetectAttentionMask pass; + pass.run_on_model(build_gemma4_masked_fill_sliding(window)); + const auto& info = pass.get_mask_info(); + EXPECT_EQ(info.mask_type, MaskInfo::MaskType::SlidingWindow); + EXPECT_EQ(info.window_size, window); +} + +// Mixed graph: one masked_fill-style sliding-window layer alongside a +// separate plain-causal SDPA layer (mirrors a real mixed sliding/full- +// attention model, e.g. Gemma-4 MoE). SlidingWindow detection must win +// regardless of GraphRewrite traversal order, otherwise HFA mask-skipping +// gets wrongly enabled for the SWA layers once the causal layer is scanned. +TEST(DetectAttentionMaskTest, RealPattern_Gemma4MaskedFillSlidingMixedWithCausalLayer_IsSlidingWindow) { + const int64_t window = 1024; + auto sliding_model = build_gemma4_masked_fill_sliding(window); + auto causal_model = make_sdpa_model(/*is_causal=*/true); + + ov::ResultVector results = sliding_model->get_results(); + const auto& causal_results = causal_model->get_results(); + results.insert(results.end(), causal_results.begin(), causal_results.end()); + + ov::ParameterVector params = sliding_model->get_parameters(); + const auto& causal_params = causal_model->get_parameters(); + params.insert(params.end(), causal_params.begin(), causal_params.end()); + + auto mixed_model = std::make_shared(results, params); + + DetectAttentionMask pass; + pass.run_on_model(mixed_model); + const auto& info = pass.get_mask_info(); + EXPECT_EQ(info.mask_type, MaskInfo::MaskType::SlidingWindow); + EXPECT_EQ(info.window_size, window); +} + // ============================================================================ // Unknown / unmasked — must report MaskType::Unknown // ============================================================================ // Full attention — SDPA with no mask and is_causal=false. TEST(DetectAttentionMaskTest, FullAttentionSDPA_IsUnknown) { - EXPECT_EQ(detect(make_sdpa_model(/*is_causal=*/false)), - MaskInfo::MaskType::Unknown); + EXPECT_EQ(detect(make_sdpa_model(/*is_causal=*/false)), MaskInfo::MaskType::Unknown); +} + +// ============================================================================ +// Per-SDPA annotation pass — must write rt_info and expose detected mask types +// ============================================================================ + +TEST(DetectAttentionMaskTest, AnnotatePerSDPAMaskType_Causal_WritesRtInfoAndGetter) { + auto [model, add] = make_decomposed_sdpa_with_mask(/*sliding_mask=*/false); + ASSERT_NE(model, nullptr); + ASSERT_NE(add, nullptr); + + AnnotatePerSDPAMaskType pass; + pass.run_on_model(model); + + const auto& annotations = pass.get_annotations(); + ASSERT_EQ(annotations.size(), 1u); + EXPECT_EQ(annotations[0].mask_type, MaskType::Causal); + + const auto mask_types = pass.get_mask_types(); + ASSERT_EQ(mask_types.size(), 1u); + EXPECT_EQ(mask_types[0], MaskType::Causal); + + const auto& rt_info = add->get_rt_info(); + const auto it = rt_info.find(ov::npuw::NPUW_SDPA_MASK_TYPE_RT_KEY); + ASSERT_NE(it, rt_info.end()); + EXPECT_EQ(static_cast(it->second.as()), MaskType::Causal); +} + +TEST(DetectAttentionMaskTest, AnnotatePerSDPAMaskType_SlidingWindow_WritesRtInfoAndGetter) { + auto [model, add] = make_decomposed_sdpa_with_mask(/*sliding_mask=*/true); + ASSERT_NE(model, nullptr); + ASSERT_NE(add, nullptr); + + AnnotatePerSDPAMaskType pass; + pass.run_on_model(model); + + const auto& annotations = pass.get_annotations(); + ASSERT_EQ(annotations.size(), 1u); + EXPECT_EQ(annotations[0].mask_type, MaskType::SlidingWindow); + + const auto mask_types = pass.get_mask_types(); + ASSERT_EQ(mask_types.size(), 1u); + EXPECT_EQ(mask_types[0], MaskType::SlidingWindow); + + const auto& rt_info = add->get_rt_info(); + const auto it = rt_info.find(ov::npuw::NPUW_SDPA_MASK_TYPE_RT_KEY); + ASSERT_NE(it, rt_info.end()); + EXPECT_EQ(static_cast(it->second.as()), MaskType::SlidingWindow); +} + +TEST(DetectAttentionMaskTest, AnnotatePerSDPAMaskType_ClearsResultsBetweenRuns) { + auto [model_with_pattern, _] = make_decomposed_sdpa_with_mask(/*sliding_mask=*/true); + ASSERT_NE(model_with_pattern, nullptr); + + AnnotatePerSDPAMaskType pass; + pass.run_on_model(model_with_pattern); + ASSERT_EQ(pass.get_annotations().size(), 1u); + + auto model_without_pattern = make_sdpa_model(/*is_causal=*/false); + ASSERT_NE(model_without_pattern, nullptr); + pass.run_on_model(model_without_pattern); + EXPECT_TRUE(pass.get_annotations().empty()); + EXPECT_TRUE(pass.get_mask_types().empty()); +} + +TEST(DetectAttentionMaskTest, AnnotatePerSDPAMaskType_MixedSWAAndGlobal_ReturnsBothMaskTypes) { + auto mixed = make_mixed_decomposed_sdpa_model(); + ASSERT_NE(mixed.model, nullptr); + ASSERT_NE(mixed.add_global, nullptr); + ASSERT_NE(mixed.add_swa, nullptr); + + AnnotatePerSDPAMaskType pass; + pass.run_on_model(mixed.model); + + const auto mask_types = pass.get_mask_types(); + ASSERT_EQ(mask_types.size(), 2u); + std::multiset type_set(mask_types.begin(), mask_types.end()); + EXPECT_EQ(type_set.count(MaskType::Causal), 1u); + EXPECT_EQ(type_set.count(MaskType::SlidingWindow), 1u); + + const auto& annotations = pass.get_annotations(); + ASSERT_EQ(annotations.size(), 2u); + std::multiset annotation_types; + for (const auto& annotation : annotations) + annotation_types.insert(annotation.mask_type); + EXPECT_EQ(annotation_types.count(MaskType::Causal), 1u); + EXPECT_EQ(annotation_types.count(MaskType::SlidingWindow), 1u); +} + +TEST(DetectAttentionMaskTest, AnnotatePerSDPAMaskType_MixedSWAAndGlobal_WritesCorrectRtInfo) { + auto mixed = make_mixed_decomposed_sdpa_model(); + ASSERT_NE(mixed.model, nullptr); + + AnnotatePerSDPAMaskType pass; + pass.run_on_model(mixed.model); + + auto get_mask_type_from_rt_info = [](const std::shared_ptr& add_node) { + const auto& rt_info = add_node->get_rt_info(); + const auto it = rt_info.find(ov::npuw::NPUW_SDPA_MASK_TYPE_RT_KEY); + EXPECT_NE(it, rt_info.end()); + return static_cast(it->second.as()); + }; + + EXPECT_EQ(get_mask_type_from_rt_info(mixed.add_global), MaskType::Causal); + EXPECT_EQ(get_mask_type_from_rt_info(mixed.add_swa), MaskType::SlidingWindow); } diff --git a/src/plugins/intel_npu/tests/unit/npuw/write_kv_slice_sliding_test.cpp b/src/plugins/intel_npu/tests/unit/npuw/write_kv_slice_sliding_test.cpp new file mode 100644 index 00000000000000..b0f7b315418072 --- /dev/null +++ b/src/plugins/intel_npu/tests/unit/npuw/write_kv_slice_sliding_test.cpp @@ -0,0 +1,205 @@ +// Copyright (C) 2018-2026 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#include + +#include +#include + +#include "infer_request_utils.hpp" +#include "openvino/openvino.hpp" +#include "openvino/runtime/make_tensor.hpp" + +namespace ov::test::npuw { + +namespace { + +namespace uu = ov::npuw::util; + +ov::SoPtr make_cpu_tensor(const ov::Shape& shape) { + return ov::get_tensor_impl(ov::Tensor(ov::element::f32, shape)); +} + +// [1, heads, seq_len, emb] for kv_dim == 2, [1, heads, emb, seq_len] for kv_dim == 3. +ov::Shape kv_shape(uint32_t kv_dim, uint32_t seq_len, uint32_t heads = 2u, uint32_t emb = 3u) { + return (kv_dim == 3u) ? ov::Shape{1u, heads, emb, seq_len} : ov::Shape{1u, heads, seq_len, emb}; +} + +// Generic element-wise walk over a (possibly non-contiguous / ROI) f32 tensor, visiting +// every logical element in row-major order. Works regardless of kv_dim/layout, so it can +// fill or read back a single-token slice (which is only contiguous for some layouts). +template +void for_each_element(const ov::SoPtr& tensor, Fn&& fn) { + const auto& shape = tensor->get_shape(); + const auto& strides = tensor->get_strides(); // byte strides + auto* base = static_cast(tensor->data()); + std::vector idx(shape.size(), 0u); + const size_t total = tensor->get_size(); + for (size_t linear = 0; linear < total; ++linear) { + size_t byte_offset = 0u; + for (size_t d = 0; d < shape.size(); ++d) { + byte_offset += idx[d] * strides[d]; + } + fn(*reinterpret_cast(base + byte_offset)); + for (int d = static_cast(shape.size()) - 1; d >= 0; --d) { + if (++idx[d] < shape[d]) { + break; + } + idx[d] = 0u; + } + } +} + +// Fills physical column `pos` along `kv_dim` with a single repeated value, so the token's +// identity can be recovered later via a single read (and cross-checked via +// expect_token_value below). +void write_token_value(const ov::SoPtr& tensor, uint32_t kv_dim, uint32_t pos, float value) { + auto slice = uu::make_tensor_slice(tensor, kv_dim, pos, pos + 1); + for_each_element(slice, [&](float& v) { + v = value; + }); +} + +// Asserts every element of physical column `pos` along `kv_dim` equals `value`. +void expect_token_value(const ov::SoPtr& tensor, uint32_t kv_dim, uint32_t pos, float value) { + auto slice = uu::make_tensor_slice(tensor, kv_dim, pos, pos + 1); + for_each_element(slice, [&](float& v) { + EXPECT_FLOAT_EQ(v, value) << "kv_dim=" << kv_dim << " pos=" << pos; + }); +} + +// Builds a source tensor of `count` tokens along `kv_dim`, whose logical token `i` (i.e. +// physical column i, since a freshly-produced source is always contiguous/right-aligned) +// is filled with `first_value + i`. +ov::SoPtr make_src_tokens(uint32_t kv_dim, uint32_t count, float first_value) { + auto src = make_cpu_tensor(kv_shape(kv_dim, count)); + for (uint32_t i = 0; i < count; ++i) { + write_token_value(src, kv_dim, i, first_value + static_cast(i)); + } + return src; +} + +} // namespace + +class WriteKvSliceSlidingTest : public ::testing::TestWithParam {}; + +// kv_dim is parameterized: 2 (standard [1,H,S,E] layout) and 3 (transposed-V layout). +INSTANTIATE_TEST_SUITE_P(KvDims, WriteKvSliceSlidingTest, ::testing::Values(2u, 3u)); + +TEST_P(WriteKvSliceSlidingTest, CircularWarmupMatchesLeftAligned) { + const uint32_t kv_dim = GetParam(); + const uint32_t capacity = 8u; + + auto dst_left = make_cpu_tensor(kv_shape(kv_dim, capacity)); + auto dst_circ = make_cpu_tensor(kv_shape(kv_dim, capacity)); + auto src = make_src_tokens(kv_dim, 5u, /*first_value=*/100.f); + + uu::write_kv_slice_sliding(dst_left, src, kv_dim, kv_dim, /*num_stored_tokens_before=*/0u, /*num_new_tokens=*/5u, + uu::SlidingBufferLayout::LeftAligned); + uu::write_kv_slice_sliding(dst_circ, src, kv_dim, kv_dim, /*num_stored_tokens_before=*/0u, /*num_new_tokens=*/5u, + uu::SlidingBufferLayout::Circular); + + for (uint32_t i = 0; i < 5u; ++i) { + expect_token_value(dst_left, kv_dim, i, 100.f + static_cast(i)); + expect_token_value(dst_circ, kv_dim, i, 100.f + static_cast(i)); + } +} + +TEST_P(WriteKvSliceSlidingTest, CircularWrapsToStartWithoutSplit) { + const uint32_t kv_dim = GetParam(); + const uint32_t capacity = 8u; + + auto dst = make_cpu_tensor(kv_shape(kv_dim, capacity)); + // Pretend absolute positions [0..7] are already stored: physical slot i holds + // value (1000 + i). + for (uint32_t i = 0; i < capacity; ++i) { + write_token_value(dst, kv_dim, i, 1000.f + static_cast(i)); + } + + // One new token arrives at absolute position 8 -> physical slot (8 % 8) == 0. + auto src = make_src_tokens(kv_dim, 1u, /*first_value=*/2000.f); + uu::write_kv_slice_sliding(dst, src, kv_dim, kv_dim, /*num_stored_tokens_before=*/capacity, /*num_new_tokens=*/1u, + uu::SlidingBufferLayout::Circular); + + expect_token_value(dst, kv_dim, 0u, 2000.f); + for (uint32_t i = 1; i < capacity; ++i) { + expect_token_value(dst, kv_dim, i, 1000.f + static_cast(i)); + } +} + +TEST_P(WriteKvSliceSlidingTest, CircularWriteSplitsAcrossWrapBoundary) { + const uint32_t kv_dim = GetParam(); + const uint32_t capacity = 8u; + + auto dst = make_cpu_tensor(kv_shape(kv_dim, capacity)); + // Physical layout represents logical positions [8,9,10,11,12,5,6,7] at physical + // slots [0,1,2,3,4,5,6,7] respectively (i.e. 13 tokens (0..12) already written). + const std::vector initial = {2000.f, 2001.f, 2002.f, 2003.f, 2004.f, 1005.f, 1006.f, 1007.f}; + for (uint32_t i = 0; i < capacity; ++i) { + write_token_value(dst, kv_dim, i, initial[i]); + } + + // 6 new tokens arrive in one call, absolute positions [13..18]. + // dst_start = 13 % 8 = 5, tokens_to_write = 6 -> wraps (5 + 6 > 8): + // leg 1: physical [5,8) <- src tokens [0,1,2] (positions 13,14,15) + // leg 2: physical [0,3) <- src tokens [3,4,5] (positions 16,17,18) + auto src = make_src_tokens(kv_dim, 6u, /*first_value=*/3013.f); + uu::write_kv_slice_sliding(dst, src, kv_dim, kv_dim, /*num_stored_tokens_before=*/13u, /*num_new_tokens=*/6u, + uu::SlidingBufferLayout::Circular); + + const std::vector expected = {3016.f, 3017.f, 3018.f, 2003.f, 2004.f, 3013.f, 3014.f, 3015.f}; + for (uint32_t i = 0; i < capacity; ++i) { + expect_token_value(dst, kv_dim, i, expected[i]); + } +} + +TEST_P(WriteKvSliceSlidingTest, CircularChunkLargerThanCapacityKeepsOnlyNewestTail) { + const uint32_t kv_dim = GetParam(); + const uint32_t capacity = 4u; + + auto dst = make_cpu_tensor(kv_shape(kv_dim, capacity)); + // A single chunked-prefill-style call writes 10 tokens (absolute positions 0..9) + // against an empty, capacity-4 buffer in one shot. Only the newest 4 (positions + // 6,7,8,9) survive; dst_start = (0 + (10-4)) % 4 = 2, wraps (2+4>4): + // leg 1: physical [2,4) <- src tokens [0,1] (positions 6,7) + // leg 2: physical [0,2) <- src tokens [2,3] (positions 8,9) + auto src = make_src_tokens(kv_dim, 10u, /*first_value=*/0.f); + uu::write_kv_slice_sliding(dst, src, kv_dim, kv_dim, /*num_stored_tokens_before=*/0u, /*num_new_tokens=*/10u, + uu::SlidingBufferLayout::Circular); + + const std::vector expected = {8.f, 9.f, 6.f, 7.f}; + for (uint32_t i = 0; i < capacity; ++i) { + expect_token_value(dst, kv_dim, i, expected[i]); + } +} + +TEST_P(WriteKvSliceSlidingTest, CircularAndLeftAlignedHoldSameLogicalContentOverManySteps) { + const uint32_t kv_dim = GetParam(); + const uint32_t capacity = 5u; + + auto dst_left = make_cpu_tensor(kv_shape(kv_dim, capacity)); + auto dst_circ = make_cpu_tensor(kv_shape(kv_dim, capacity)); + + uint32_t stored = 0u; + for (uint32_t step = 0; step < 20u; ++step) { + // One new token per step, value == its own absolute position (so recovering it + // later is a direct identity check). + auto src = make_src_tokens(kv_dim, 1u, /*first_value=*/static_cast(stored)); + uu::write_kv_slice_sliding(dst_left, src, kv_dim, kv_dim, stored, 1u, uu::SlidingBufferLayout::LeftAligned); + uu::write_kv_slice_sliding(dst_circ, src, kv_dim, kv_dim, stored, 1u, uu::SlidingBufferLayout::Circular); + stored += 1u; + + const uint32_t valid = std::min(stored, capacity); + const uint32_t window_start = stored - valid; // oldest absolute position still in window + for (uint32_t rank = 0; rank < valid; ++rank) { + const uint32_t abs_pos = window_start + rank; + // LeftAligned: rank-th oldest surviving token sits at physical index `rank`. + expect_token_value(dst_left, kv_dim, rank, static_cast(abs_pos)); + // Circular: token at abs_pos always sits at physical index (abs_pos % capacity). + expect_token_value(dst_circ, kv_dim, abs_pos % capacity, static_cast(abs_pos)); + } + } +} + +} // namespace ov::test::npuw