Skip to content

[NPUW] Replacing longrope pattern with precalculated values #33011

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
dmatveev merged 7 commits into from
Mar 4, 2026
Merged

[NPUW] Replacing longrope pattern with precalculated values #33011

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
2a01767
[NPUW] Replacing longrope pattern with precalculated values
AlexanderKalistratov Nov 25, 2025
a1e9ab1
Remove stractural bindings
AlexanderKalistratov Nov 25, 2025
6862daa
Merge branch 'master' into npu_longrope_support
AlexanderKalistratov Jan 31, 2026
bfbd69e
Merge branch 'master' into npu_longrope_support
AlexanderKalistratov Feb 10, 2026
f2c3a80
Merge branch 'master' into npu_longrope_support
AlexanderKalistratov Feb 24, 2026
1a2963b
Fix copyright
AlexanderKalistratov Feb 24, 2026
3a23df4
fix copyright
AlexanderKalistratov Feb 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 6 additions & 2 deletions src/plugins/intel_npu/src/plugin/npuw/llm_compiled_model.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1952,7 +1952,9 @@ ov::npuw::LLMCompiledModel::LLMCompiledModel(const std::shared_ptr<ov::Model>& m

if (!is_best || (max_prompt_len >= CACHE_ROPE_START)) {
LOG_DEBUG("Enable RoPE Cache for prefill");
ov::npuw::patterns::pre_compute::RopeCache rope_prefill_cacher(max_prompt_len);
ov::npuw::patterns::pre_compute::RopeCache rope_prefill_cacher(
max_prompt_len,
ov::npuw::LLMInferRequest::layer_names::longrope_input);
rope_prefill_cacher.run_on_model(prefill_model);
}

Expand All @@ -1961,7 +1963,9 @@ ov::npuw::LLMCompiledModel::LLMCompiledModel(const std::shared_ptr<ov::Model>& m
const uint32_t kv_size = m_kvcache_sizes[i];
if (!is_best || (kv_size >= CACHE_ROPE_START)) {
LOG_DEBUG("Enable RoPE Cache for generate variant with size: " << kv_size);
ov::npuw::patterns::pre_compute::RopeCache rope_cacher(kv_size);
ov::npuw::patterns::pre_compute::RopeCache rope_cacher(
kv_size,
ov::npuw::LLMInferRequest::layer_names::longrope_input);
rope_cacher.run_on_model(generate_model_variants[i]);
}
}
Expand Down
10 changes: 6 additions & 4 deletions src/plugins/intel_npu/src/plugin/npuw/llm_infer_base_request.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
// Copyright (C) 2025 Intel Corporation
// Copyright (C) 2018-2026 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//

Expand All @@ -22,14 +22,16 @@ class LLMInferBaseRequest : public ov::ISyncInferRequest {
static constexpr const char* output_embeds = "npuw_output_embed";
static constexpr const char* logits = "logits";
static constexpr const char* token_type_ids = "token_type_ids";
static constexpr const char* gemma_sliding_mask = "npuw_gemma_sliding_mask";
static constexpr const char* longrope_input = "npuw_longrope_input";
};

struct layer_ids {
static constexpr uint32_t INPUT_IDS_SEQ_LEN_DIM = 1;
static constexpr std::size_t kStartOutputKVCacheLayers = 1;
};

using PortsMap = std::unordered_map<std::string, ov::Output<const ov::Node>>;

explicit LLMInferBaseRequest(const std::shared_ptr<LLMCompiledModel>& compiled_model)
: ISyncInferRequest(compiled_model),
m_npuw_llm_compiled_model(compiled_model) {}
Expand All @@ -44,8 +46,8 @@ class LLMInferBaseRequest : public ov::ISyncInferRequest {

protected:
void update_kvcache_for(std::shared_ptr<ov::IAsyncInferRequest> request,
const std::unordered_map<std::string, ov::Output<const ov::Node>>& in_ports,
const std::unordered_map<std::string, ov::Output<const ov::Node>>& out_ports,
const PortsMap& in_ports,
const PortsMap& out_ports,
uint32_t num_tokens,
bool v_transposed);
void init_tensor(const ov::Output<const ov::Node>& port);
Expand Down
19 changes: 19 additions & 0 deletions src/plugins/intel_npu/src/plugin/npuw/llm_infer_request.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
#include "util.hpp"

namespace {
using ov::npuw::LLMInferRequest;

void copy_columns_by_row_chunks_2d(ov::SoPtr<ov::ITensor> src, ov::SoPtr<ov::ITensor>& dst) {
const auto& src_shape = src->get_shape();
Expand Down Expand Up @@ -88,6 +89,21 @@ std::pair<uint32_t, uint32_t> get_lora_dims_by_name(const std::string& state_nam
return std::make_pair(low_rank_dim, full_rank_dim);
}

void process_longrope(const std::shared_ptr<ov::IAsyncInferRequest>& infer_req,
const LLMInferRequest::PortsMap& ports,
const ov::SoPtr<ov::ITensor>& position_ids) {
if (auto longrope_port_it = ports.find(LLMInferRequest::layer_names::longrope_input);
longrope_port_it != ports.end()) {
auto* pos_ids_data = position_ids->data<int64_t>();
// assuming position_ids are constantly non-deacreasing.
// this potentially could be not true. Alternative is to find max value in position_ids
auto max_pos_id = pos_ids_data[position_ids->get_size() - 1];
Copy link
Contributor Author

@AlexanderKalistratov AlexanderKalistratov Feb 24, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Check if attention mask is right padded


auto longrope_input = infer_req->get_tensor(longrope_port_it->second);
longrope_input->data<int64_t>()[0] = max_pos_id;
}
}

} // anonymous namespace

void ov::npuw::LLMInferRequest::init_lora_states() {
Expand Down Expand Up @@ -811,6 +827,7 @@ void ov::npuw::LLMInferRequest::infer_prefill(ov::SoPtr<ov::ITensor> input_ids,

prepare_for_new_conversation(prompt_length);

process_longrope(m_prefill_request, m_prefill_in_ports, position_ids);
const bool use_chunk_prefill = m_npuw_llm_compiled_model->m_use_chunk_prefill;
if (use_chunk_prefill) {
OPENVINO_ASSERT(!token_type_ids,
Expand Down Expand Up @@ -880,6 +897,8 @@ void ov::npuw::LLMInferRequest::infer_generate(ov::SoPtr<ov::ITensor> input_ids,
OPENVINO_THROW("KV-Cache is full.");
}

process_longrope(m_kvcache_request, m_kvcache_in_ports, position_ids);

// FIXME: these tensors should be shared between the parent & child models
// NB: input_ids can be either fp32(VLM) or i64(LLM)
auto kv_input_ids = m_kvcache_request->get_tensor(m_kvcache_in_ports.at(m_input_ids_name));
Expand Down
16 changes: 6 additions & 10 deletions src/plugins/intel_npu/src/plugin/npuw/llm_infer_request.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -78,21 +78,17 @@ class LLMInferRequest : public ov::npuw::LLMInferBaseRequest {
std::shared_ptr<ov::IAsyncInferRequest> m_lm_head_request;
ov::SoPtr<ov::ITensor> m_logits;

std::unordered_map<std::string, ov::Output<const ov::Node>> m_prefill_in_ports;
std::unordered_map<std::string, ov::Output<const ov::Node>> m_prefill_out_ports;
PortsMap m_prefill_in_ports;
PortsMap m_prefill_out_ports;

// Ports for the currently selected generate model variant (set once per conversation in
// prepare_for_new_conversation)
std::unordered_map<std::string, ov::Output<const ov::Node>> m_kvcache_in_ports;
std::unordered_map<std::string, ov::Output<const ov::Node>> m_kvcache_out_ports;
PortsMap m_kvcache_in_ports;
PortsMap m_kvcache_out_ports;

// Ports for all generate model variants - maps from request pointer to its input/output ports
std::unordered_map<std::shared_ptr<ov::IAsyncInferRequest>,
std::unordered_map<std::string, ov::Output<const ov::Node>>>
m_generate_variant_in_ports;
std::unordered_map<std::shared_ptr<ov::IAsyncInferRequest>,
std::unordered_map<std::string, ov::Output<const ov::Node>>>
m_generate_variant_out_ports;
std::unordered_map<std::shared_ptr<ov::IAsyncInferRequest>, PortsMap> m_generate_variant_in_ports;
std::unordered_map<std::shared_ptr<ov::IAsyncInferRequest>, PortsMap> m_generate_variant_out_ports;

ov::Output<const ov::Node> m_lm_head_logits_port;

Expand Down
173 changes: 137 additions & 36 deletions src/plugins/intel_npu/src/plugin/npuw/partitioning/patterns/pre_compute.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
#include "transformations/common_optimizations/fuse_rotary_positional_embeddings.hpp"

namespace opp = ov::pass::pattern;
namespace pre_compute = ov::npuw::patterns::pre_compute;

namespace {
// TODO: copied from common tests
Expand Down Expand Up @@ -49,6 +50,46 @@ static ov::OutputVector makeCosSinCache(const size_t max_position_embeddings,

return {Cos, Sin};
}

void replaceSinCosByCache(int max_prompt_len, const ov::OutputVector& cache, const pre_compute::RopePatternDesc* rpe) {
Copy link
Contributor

@AsyaPronina AsyaPronina Mar 4, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Might be to align with makeCosSinCache we need to swap Sin and Cos in name by places?

auto inv_freq_size = ov::shape_size(rpe->matched_inv_freq->get_shape());

LOG_VERB("Making sin-cos cache of size: " << max_prompt_len << "x" << inv_freq_size);

// Step 1: Define axis (gather along axis 1)
auto axis = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{1}, {1});

// Step 2: Apply Gather for cos and sin
auto gather_cos = std::make_shared<ov::op::v8::Gather>(cache[0], rpe->matched_position_ids, axis);
auto gather_sin = std::make_shared<ov::op::v8::Gather>(cache[1], rpe->matched_position_ids, axis);
LOG_VERB("Created gather op facilitate LUT search: " << gather_cos->get_name() << ", " << gather_cos->get_shape());
Copy link
Contributor

@AsyaPronina AsyaPronina Mar 4, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we need verbose for the sin also?


// Step 2: convert fp16->fp32
auto cos_fp32 = std::make_shared<ov::op::v0::Convert>(gather_cos, ov::element::f32);
auto sin_fp32 = std::make_shared<ov::op::v0::Convert>(gather_sin, ov::element::f32);

// Create the squeeze operation required after gather
auto squeeze_cos = std::make_shared<ov::op::v0::Squeeze>(cos_fp32, axis);
auto squeeze_sin = std::make_shared<ov::op::v0::Squeeze>(sin_fp32, axis);

LOG_VERB("Created squeeze_cos op to reduce axis=1: " << squeeze_cos->get_name() << ", "
<< squeeze_cos->get_shape());
LOG_VERB("Created squeeze_sin op to reduce axis=1: " << squeeze_sin->get_name() << ", "
<< squeeze_sin->get_shape());

LOG_VERB("Rope cos detected at: " << rpe->matched_cos->get_name() << ", replacing by cache node: "
<< gather_cos->get_name() << ", " << gather_cos->get_shape());
LOG_VERB("Rope sin detected at: " << rpe->matched_sin->get_name() << ", replacing by cache node: "
<< gather_sin->get_name() << ", " << gather_sin->get_shape());

// replacing sin with gather op
Copy link
Contributor

@AsyaPronina AsyaPronina Mar 4, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

cos and sin

ov::replace_node(rpe->matched_cos, squeeze_cos);
ov::replace_node(rpe->matched_sin, squeeze_sin);

// disconnecting gather from rest or subgraph started from concat_1
Copy link
Contributor

@AsyaPronina AsyaPronina Mar 4, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Minor comment: typo in of

auto gather_input_to_concat = rpe->matched_concat->input(0);
gather_input_to_concat.get_source_output().remove_target_input(gather_input_to_concat);
}
} // namespace

ov::npuw::patterns::pre_compute::RopePatternLLama2::RopePatternLLama2() : matcher("sin-cos-matcher") {
Expand Down Expand Up @@ -87,54 +128,114 @@ ov::npuw::patterns::pre_compute::RopePatternLLama2::RopePatternLLama2() : matche
matcher.register_patterns({output_sin, output_cos}, make_matcher_callback());
}

ov::npuw::patterns::pre_compute::RopeCacheMatcher::RopeCacheMatcher(const uint32_t max_prompt_len,
const std::shared_ptr<ov::Model>& model) {
auto rpe = std::make_shared<RopePatternLLama2>();
ov::npuw::patterns::pre_compute::LongRopePatternPhi::LongRopePatternPhi() : matcher("sin-cos-matcher") {
auto MakeConstant = []() {
Copy link
Contributor Author

@AlexanderKalistratov AlexanderKalistratov Feb 24, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Run pass conditionally

return opp::wrap_type<ov::op::v0::Constant>();
};

rpe->transform_cb = [&]() {
auto inv_freq_size = ov::shape_size(rpe->matched_inv_freq->get_shape());
auto make_select_pattern = [&](const std::shared_ptr<ov::Node>& position_ids,
const std::shared_ptr<ov::Node>& inv_freq_short,
const std::shared_ptr<ov::Node>& inv_freq_long) {
auto red_max = opp::wrap_type<ov::op::v1::ReduceMax>({position_ids, MakeConstant()});
auto add = opp::wrap_type<ov::op::v1::Add>({red_max, MakeConstant()});
// max(position_ids) + 1 <= original_max_position_embeddings
auto leq = opp::wrap_type<ov::op::v1::LessEqual>({add, MakeConstant()});

LOG_VERB("Making sin-cos cache of size: " << max_prompt_len << "x" << inv_freq_size);
auto inv_freq_short_conv = opp::optional<ov::op::v0::Convert>({inv_freq_short->output(0)});
auto inv_freq_long_conv = opp::optional<ov::op::v0::Convert>({inv_freq_long->output(0)});

// shapes that matches max possible position
auto cache = makeCosSinCache(max_prompt_len, rpe->matched_inv_freq);
// max(position_ids) + 1 <= original_max_position_embeddings ? short_factor : long_factor;
auto select = opp::wrap_type<ov::op::v1::Select>({leq, inv_freq_short_conv, inv_freq_long_conv});
auto unsqueeze = opp::optional<ov::op::v0::Unsqueeze>({select, MakeConstant()});
auto unsqueeze_1 = opp::optional<ov::op::v0::Unsqueeze>({unsqueeze, MakeConstant()});

return std::make_tuple(unsqueeze_1, leq, red_max);
};

auto position_ids = opp::wrap_type<ov::op::v0::Parameter>();

auto inv_freq_short = MakeConstant();
auto inv_freq_long = MakeConstant();

auto select_cond_max_pos_id = make_select_pattern(position_ids, inv_freq_short, inv_freq_long);
auto select = std::get<0>(select_cond_max_pos_id);
auto cond = std::get<1>(select_cond_max_pos_id);
auto max_pos_id = std::get<2>(select_cond_max_pos_id);

auto shape_of = opp::wrap_type<ov::op::v3::ShapeOf>({opp::any_input()});
auto gather = opp::wrap_type<ov::op::v8::Gather>({shape_of, opp::any_input(), opp::any_input()});
auto concat_1 = opp::wrap_type<ov::op::v0::Concat>({gather, opp::any_input(), opp::any_input()});
// here we can seen inverse frequencies as a parameter or constant depending on partitioner passes
auto broadcast = opp::wrap_type<ov::op::v3::Broadcast>({select, concat_1});
auto unsqueeze = opp::wrap_type<ov::op::v0::Unsqueeze>({position_ids, MakeConstant()});
auto convert = opp::wrap_type<ov::op::v0::Convert>({unsqueeze});
auto matmul = opp::wrap_type<ov::op::v0::MatMul>({broadcast, convert});
auto transpose = opp::wrap_type<ov::op::v1::Transpose>({matmul, opp::any_input()});
auto concat_2 = opp::wrap_type<ov::op::v0::Concat>({transpose, opp::any_input()});
auto output_sin = opp::wrap_type<ov::op::v0::Sin>({concat_2});
auto output_cos = opp::wrap_type<ov::op::v0::Cos>({concat_2});

// Step 1: Define axis (gather along axis 1)
auto axis = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{1}, {1});
init_cb = [=](const auto& matches) {
const auto& map_sin = matches.at(output_sin)[0];
const auto& map_cos = matches.at(output_cos)[0];

// Step 2: Apply Gather for cos and sin
auto gather_cos = std::make_shared<ov::op::v8::Gather>(cache[0], rpe->matched_position_ids, axis);
auto gather_sin = std::make_shared<ov::op::v8::Gather>(cache[1], rpe->matched_position_ids, axis);
LOG_VERB("Created gather op facilitate LUT search: " << gather_cos->get_name() << ", "
<< gather_cos->get_shape());
this->matched_position_ids = map_sin.at(position_ids).get_node_shared_ptr();
this->matched_concat = map_sin.at(concat_1).get_node_shared_ptr();
this->matched_inv_freq = map_sin.at(inv_freq_short).get_node_shared_ptr();
this->matched_inv_freq_long = map_sin.at(inv_freq_long).get_node_shared_ptr();
this->matched_cond = map_sin.at(cond).get_node_shared_ptr();
this->max_pos_id = map_sin.at(max_pos_id).get_node_shared_ptr();

// Step 2: convert fp16->fp32
auto cos_fp32 = std::make_shared<ov::op::v0::Convert>(gather_cos, ov::element::f32);
auto sin_fp32 = std::make_shared<ov::op::v0::Convert>(gather_sin, ov::element::f32);
this->matched_cos = map_cos.at(output_cos).get_node_shared_ptr();
this->matched_sin = map_sin.at(output_sin).get_node_shared_ptr();

// Create the squeeze operation required after gather
auto squeeze_cos = std::make_shared<ov::op::v0::Squeeze>(cos_fp32, axis);
auto squeeze_sin = std::make_shared<ov::op::v0::Squeeze>(sin_fp32, axis);
LOG_VERB("Rope found : sin=" << matched_sin->get_name() << ", cos=" << matched_cos->get_name());

LOG_VERB("Created squeeze_cos op to reduce axis=1: " << squeeze_cos->get_name() << ", "
<< squeeze_cos->get_shape());
LOG_VERB("Created squeeze_sin op to reduce axis=1: " << squeeze_sin->get_name() << ", "
<< squeeze_sin->get_shape());
return true;
};

LOG_VERB("Rope cos detected at: " << rpe->matched_cos->get_name() << ", replacing by cache node: "
<< gather_cos->get_name() << ", " << gather_cos->get_shape());
LOG_VERB("Rope sin detected at: " << rpe->matched_sin->get_name() << ", replacing by cache node: "
<< gather_sin->get_name() << ", " << gather_sin->get_shape());
matcher.register_patterns({output_sin, output_cos}, make_matcher_callback());
}

// replacing sin with gather op
ov::replace_node(rpe->matched_cos, squeeze_cos);
ov::replace_node(rpe->matched_sin, squeeze_sin);
ov::npuw::patterns::pre_compute::RopeCacheMatcher::RopeCacheMatcher(const uint32_t max_prompt_len,
const std::shared_ptr<ov::Model>& model,
const std::string& longrope_input_name) {
auto rpe = std::make_shared<RopePatternLLama2>();

// disconnecting gather from rest or subgraph started from concat_1
auto gather_input_to_concat = rpe->matched_concat->input(0);
gather_input_to_concat.get_source_output().remove_target_input(gather_input_to_concat);
rpe->transform_cb = [&]() {
auto cache = makeCosSinCache(max_prompt_len, rpe->matched_inv_freq);
replaceSinCosByCache(max_prompt_len, cache, rpe.get());
};
rpe->run_on_model(model);

auto long_rpe = std::make_shared<LongRopePatternPhi>();

std::shared_ptr<ov::op::v0::Parameter> max_pos_id_param;
long_rpe->transform_cb = [&]() {
auto cache_short = makeCosSinCache(max_prompt_len, long_rpe->matched_inv_freq);
auto cache_long = makeCosSinCache(max_prompt_len, long_rpe->matched_inv_freq_long);

auto select_cos = std::make_shared<ov::op::v1::Select>(long_rpe->matched_cond, cache_short[0], cache_long[0]);
auto select_sin = std::make_shared<ov::op::v1::Select>(long_rpe->matched_cond, cache_short[1], cache_long[1]);

replaceSinCosByCache(max_prompt_len, {select_cos, select_sin}, long_rpe.get());

auto max_pos_id_out = long_rpe->max_pos_id->output(0);
max_pos_id_param.reset(new ov::op::v0::Parameter(max_pos_id_out.get_element_type(), {1}));
max_pos_id_param->set_friendly_name(longrope_input_name);
max_pos_id_out.replace(max_pos_id_param->output(0));
};
long_rpe->run_on_model(model);

if (max_pos_id_param) {
model->add_parameters({max_pos_id_param});
for (auto&& input : model->inputs()) {
if (input.get_node() == max_pos_id_param.get()) {
input.set_names({max_pos_id_param->get_friendly_name()});
}
}
}
model->validate_nodes_and_infer_types();
}

ov::npuw::patterns::pre_compute::RopeInverseFreq::RopeInverseFreq(
Expand All @@ -154,6 +255,6 @@ ov::npuw::patterns::pre_compute::RopeInverseFreq::RopeInverseFreq(
}

bool ov::npuw::patterns::pre_compute::RopeCache::run_on_model(const std::shared_ptr<ov::Model>& model) {
ov::npuw::patterns::pre_compute::RopeCacheMatcher ropeCache(m_max_prompt_len, model);
ov::npuw::patterns::pre_compute::RopeCacheMatcher ropeCache(m_max_prompt_len, model, m_longrope_input_name);
return true;
}
Loading
Loading