Skip to content

Commit c99494a

Browse files
author
Evgenii Smirnov
committed
[NPUW] Support KV cache i8 compression with pyramid attention
1 parent b8a1c7e commit c99494a

10 files changed

Lines changed: 930 additions & 344 deletions

File tree

‎src/plugins/intel_npu/src/plugin/npuw/partitioning/online/snapshot.cpp‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -770,6 +770,7 @@ void Snapshot::earlyRegroup() {
770770
HNDL_ATTN(SDPA);
771771
HNDL_ATTN(SDPADecomposed);
772772
HNDL_ATTN(GQA);
773+
HNDL_ATTN(SDPACompressed);
773774
#undef HNDL_MOE
774775
#undef HNDL_ATTN
775776
#undef HNDL_FAKE

‎src/plugins/intel_npu/src/plugin/npuw/partitioning/online/utils/utils.hpp‎

Lines changed: 13 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -92,20 +92,19 @@ std::vector<Isolate> getIsolates(const std::string& isolates_unparsed);
9292
std::vector<std::string> getNoFolds(const ::intel_npu::Config& cfg);
9393
std::vector<std::string> getNoFolds(const std::string& nofolds_unparsed);
9494

95-
static const std::map<std::string, std::string> ISOL_PRESETS = {
96-
{"COMPUTE",
97-
"P:DQMatMulGQu4/compute,P:DQMatMulCWu4/compute,"
98-
"P:DQMatMulGQi4/compute,P:DQMatMulCWi4/compute,"
99-
"P:DQMatMulConv/compute,"
100-
"P:VocabMatMul/compute,"
101-
"P:RMSNorm/compute,P:RMSNorm2/compute,"
102-
"P:RMSNorm3/compute,P:RMSNorm4/compute,"
103-
"P:VariadicSplit/compute"},
104-
{"FAKE", "P:FakeConvert/fake,P:FakeQuantize/fake"},
105-
{"ATTN", "P:SDPA/attn,P:SDPADecomposed/attn,P:GQA/attn"},
106-
{"MOE",
107-
"P:GPTOSSExpert/expert,P:GPTOSSRouter/router,"
108-
"P:Qwen3Expert/expert,P:Qwen3Router/router"}};
95+
static const std::map<std::string, std::string> ISOL_PRESETS = {{"COMPUTE",
96+
"P:DQMatMulGQu4/compute,P:DQMatMulCWu4/compute,"
97+
"P:DQMatMulGQi4/compute,P:DQMatMulCWi4/compute,"
98+
"P:DQMatMulConv/compute,"
99+
"P:VocabMatMul/compute,"
100+
"P:RMSNorm/compute,P:RMSNorm2/compute,"
101+
"P:RMSNorm3/compute,P:RMSNorm4/compute,"
102+
"P:VariadicSplit/compute"},
103+
{"FAKE", "P:FakeConvert/fake,P:FakeQuantize/fake"},
104+
{"ATTN", "P:SDPA/attn,P:SDPADecomposed/attn,P:GQA/attn,P:SDPACompressed/attn"},
105+
{"MOE",
106+
"P:GPTOSSExpert/expert,P:GPTOSSRouter/router,"
107+
"P:Qwen3Expert/expert,P:Qwen3Router/router"}}
109108
} // namespace util
110109

111110
} // namespace online

‎src/plugins/intel_npu/src/plugin/npuw/partitioning/patterns/sdpa.cpp‎

Lines changed: 129 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -223,6 +223,135 @@ SDPADecomposed::SDPADecomposed(const std::shared_ptr<ov::npuw::online::Snapshot>
223223
register_matcher(std::make_shared<opp::Matcher>(reshape3, "TagSDPADecomposed"), std::move(callback));
224224
}
225225

226+
/*
227+
Decomposed SDPA Pattern with Dynamic Dequantization (i8 KV cache):
228+
229+
After ConvertKVCacheToPrecision(i8), past KV cache inputs have DQ nodes:
230+
[any_input] → Subtract(zp) → Multiply(scale) → Concat
231+
instead of:
232+
Convert → Concat
233+
234+
Full pattern:
235+
Convert
236+
|
237+
opt:Subtract (zp)
238+
|
239+
Multiply (scale)
240+
\ /
241+
Concat
242+
|
243+
opt:Unsqueeze
244+
|
245+
opt:Broadcast Convert
246+
| \ |
247+
opt:Reshape opt:Subtract (zp)
248+
| |
249+
| Multiply (scale)
250+
| \ /
251+
\ / Concat
252+
MatMul |
253+
\ / opt:Unsqueeze
254+
Add |
255+
| opt:Broadcast
256+
Softmax |
257+
\ opt:Reshape
258+
\ /
259+
MatMul
260+
|
261+
Transpose
262+
|
263+
Reshape
264+
|
265+
*/
266+
267+
SDPACompressed::SDPACompressed(const std::shared_ptr<ov::npuw::online::Snapshot>& snapshot,
268+
const std::string& isol_tag) {
269+
// Key path: opt:Convert → opt:Subtract(opt:Convert(any), any) → Multiply(any) → Concat(any)
270+
// Convert is optional to handle models where the past KV is already in the expected type.
271+
// Subtract is optional to handle both asymmetric (with zp) and symmetric (without zp) DQ.
272+
// The zp input also has an optional Convert (i8→f32) that must be isolated
273+
// to preserve the original DQ zp parameter name for PyramidAttention::from().
274+
auto convert1 = opp::optional<ov::op::v0::Convert>({opp::any_input()});
275+
auto zp_convert1 = opp::optional<ov::op::v0::Convert>({opp::any_input()});
276+
auto subtract1 = opp::optional<ov::op::v1::Subtract>({convert1, zp_convert1});
277+
auto multiply1 = opp::wrap_type<ov::op::v1::Multiply>({subtract1, opp::any_input()});
278+
auto concat1 = opp::wrap_type<ov::op::v0::Concat>({multiply1, opp::any_input()});
279+
280+
// GQA optional nodes — single consumer guard
281+
auto single_user = [](const ov::Output<ov::Node>& output) {
282+
return output.get_target_inputs().size() == 1;
283+
};
284+
auto unsqueeze1 = opp::optional<ov::op::v0::Unsqueeze>({concat1, opp::any_input()}, single_user);
285+
auto broadcast1 = opp::optional<ov::op::v3::Broadcast>({unsqueeze1, opp::any_input()}, single_user);
286+
auto reshape1 = opp::optional<ov::op::v1::Reshape>({broadcast1, opp::any_input()}, single_user);
287+
288+
// Value path: opt:Convert → opt:Subtract(opt:Convert(any), any) → Multiply(any) → Concat(any)
289+
auto convert2 = opp::optional<ov::op::v0::Convert>({opp::any_input()});
290+
auto zp_convert2 = opp::optional<ov::op::v0::Convert>({opp::any_input()});
291+
auto subtract2 = opp::optional<ov::op::v1::Subtract>({convert2, zp_convert2});
292+
auto multiply2 = opp::wrap_type<ov::op::v1::Multiply>({subtract2, opp::any_input()});
293+
auto concat2 = opp::wrap_type<ov::op::v0::Concat>({multiply2, opp::any_input()});
294+
295+
// GQA optional nodes — same single-consumer guard
296+
auto unsqueeze2 = opp::optional<ov::op::v0::Unsqueeze>({concat2, opp::any_input()}, single_user);
297+
auto broadcast2 = opp::optional<ov::op::v3::Broadcast>({unsqueeze2, opp::any_input()}, single_user);
298+
auto reshape2 = opp::optional<ov::op::v1::Reshape>({broadcast2, opp::any_input()}, single_user);
299+
300+
auto matmul1 = opp::wrap_type<ov::op::v0::MatMul>({opp::any_input(), reshape1});
301+
auto add = opp::wrap_type<ov::op::v1::Add>({matmul1, opp::any_input()});
302+
auto softmax = opp::wrap_type<ov::op::v8::Softmax>({add});
303+
304+
auto matmul2 = opp::wrap_type<ov::op::v0::MatMul>({softmax, reshape2});
305+
auto transpose = opp::wrap_type<ov::op::v1::Transpose>({matmul2, opp::any_input()});
306+
auto reshape3 = opp::wrap_type<ov::op::v1::Reshape>({transpose, opp::any_input()});
307+
308+
auto node_to_gptr = snapshot->getNodeToGroupMap();
309+
310+
auto callback = [=](ov::pass::pattern::Matcher& m) {
311+
LOG_DEBUG("Decomposed SDPA DQ pattern matched!");
312+
313+
auto& node_to_output = m.get_pattern_value_map();
314+
315+
auto isolate_matched = [&](const auto& pattern) {
316+
auto optional_node = node_to_output.find(pattern);
317+
if (optional_node != node_to_output.end()) {
318+
auto matched_node = optional_node->second.get_node_shared_ptr();
319+
node_to_gptr->at(matched_node)->isolate(isol_tag);
320+
}
321+
};
322+
323+
// Isolate all matched nodes in the pattern
324+
isolate_matched(convert1);
325+
isolate_matched(zp_convert1);
326+
isolate_matched(subtract1);
327+
isolate_matched(multiply1);
328+
isolate_matched(concat1);
329+
isolate_matched(unsqueeze1);
330+
isolate_matched(broadcast1);
331+
isolate_matched(reshape1);
332+
333+
isolate_matched(convert2);
334+
isolate_matched(zp_convert2);
335+
isolate_matched(subtract2);
336+
isolate_matched(multiply2);
337+
isolate_matched(concat2);
338+
isolate_matched(unsqueeze2);
339+
isolate_matched(broadcast2);
340+
isolate_matched(reshape2);
341+
342+
isolate_matched(matmul1);
343+
isolate_matched(add);
344+
isolate_matched(softmax);
345+
isolate_matched(matmul2);
346+
isolate_matched(transpose);
347+
isolate_matched(reshape3);
348+
349+
return false; // root hasn't changed
350+
};
351+
352+
register_matcher(std::make_shared<opp::Matcher>(reshape3, "TagSDPADecomposedDQ"), std::move(callback));
353+
}
354+
226355
} // namespace attn
227356

228357
namespace regularize {

‎src/plugins/intel_npu/src/plugin/npuw/partitioning/patterns/sdpa.hpp‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,27 @@ class SDPADecomposed : public ov::pass::MatcherPass {
5252
SDPADecomposed(const std::shared_ptr<ov::npuw::online::Snapshot>& snapshot, const std::string& isol_tag);
5353
};
5454

55+
// Matches decomposed SDPA pattern where past KV cache inputs have been converted
56+
// to integer precision (i8/u8) with dynamic dequantization nodes inserted by
57+
// ConvertKVCacheToPrecision. The dequantization chain is:
58+
// [any_input] → Subtract(zp) → Multiply(scale) → Concat
59+
// instead of the original:
60+
// Convert → Concat
61+
class SDPACompressed : public ov::pass::MatcherPass {
62+
public:
63+
OPENVINO_MATCHER_PASS_RTTI("npuw::patterns::attn::SDPACompressed");
64+
static constexpr const char* pattern_name() {
65+
return "SDPACompressed";
66+
}
67+
static constexpr const char* isolation_tag() {
68+
return "attn";
69+
}
70+
static constexpr const char* group_name() {
71+
return "attn";
72+
}
73+
SDPACompressed(const std::shared_ptr<ov::npuw::online::Snapshot>& snapshot, const std::string& isol_tag);
74+
};
75+
5576
} // namespace attn
5677

5778
namespace regularize {

‎src/plugins/intel_npu/src/plugin/npuw/pyramid_attention.cpp‎

Lines changed: 29 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ std::optional<ov::npuw::function::Attention> create_attention_from_model(
4242
attention._mask = mask_param;
4343
attention._mask_shape = mask_param->get_shape();
4444

45-
// Add past key/value inputs to attention
45+
// Add past key/value inputs and DQ scale/zp inputs to attention
4646
const auto& params = model->get_parameters();
4747
for (const auto& param : params) {
4848
const std::string param_name = param->get_friendly_name();
@@ -56,6 +56,16 @@ std::optional<ov::npuw::function::Attention> create_attention_from_model(
5656
if (dim_iter != past_value_sequence_dims.end()) {
5757
attention._inputs.push_back(ov::npuw::function::Attention::Param{param, dim_iter->second});
5858
}
59+
} else if (ov::npuw::util::isDQScaleOrZPKey(param_name)) {
60+
if (!past_key_sequence_dims.empty()) {
61+
attention._inputs.push_back(
62+
ov::npuw::function::Attention::Param{param, past_key_sequence_dims.begin()->second});
63+
}
64+
} else if (ov::npuw::util::isDQScaleOrZPValue(param_name)) {
65+
if (!past_value_sequence_dims.empty()) {
66+
attention._inputs.push_back(
67+
ov::npuw::function::Attention::Param{param, past_value_sequence_dims.begin()->second});
68+
}
5969
}
6070
}
6171

@@ -319,6 +329,24 @@ std::optional<PyramidModelResult> process_pyramid_model(const std::shared_ptr<ov
319329
LOG_WARN("No pre-analyzed sequence dimension for past value param: " << param_name);
320330
return std::nullopt;
321331
}
332+
} else if (ov::npuw::util::isDQScaleOrZPKey(param_name)) {
333+
// Handle DynamicQuantize scale/zp parameters for past key cache.
334+
// These share the same sequence dimension as the corresponding past key parameter.
335+
if (!past_key_sequence_dims.empty()) {
336+
size_t sequence_dim_idx = past_key_sequence_dims.begin()->second;
337+
new_shape[sequence_dim_idx] = current_past_length;
338+
new_shapes[param->output(0)] = new_shape;
339+
LOG_DEBUG(" DQ key param '" << param_name << "' shape: " << original_shape << " -> " << new_shape);
340+
}
341+
} else if (ov::npuw::util::isDQScaleOrZPValue(param_name)) {
342+
// Handle DynamicQuantize scale/zp parameters for past value cache.
343+
// These share the same sequence dimension as the corresponding past value parameter.
344+
if (!past_value_sequence_dims.empty()) {
345+
size_t sequence_dim_idx = past_value_sequence_dims.begin()->second;
346+
new_shape[sequence_dim_idx] = current_past_length;
347+
new_shapes[param->output(0)] = new_shape;
348+
LOG_DEBUG(" DQ value param '" << param_name << "' shape: " << original_shape << " -> " << new_shape);
349+
}
322350
}
323351
}
324352

‎src/plugins/intel_npu/src/plugin/npuw/util.cpp‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -966,6 +966,18 @@ bool ov::npuw::util::isPastValueParam(const std::string& str) {
966966
return std::regex_match(str, pattern);
967967
}
968968

969+
bool ov::npuw::util::isDQScaleOrZPKey(const std::string& str) {
970+
// Match DynamicQuantize scale/zp parameters for past key cache
971+
static const std::regex pattern(R"(DynamicQuantize/\d+/past_key_values/key/(?:scale|zp))");
972+
return std::regex_match(str, pattern);
973+
}
974+
975+
bool ov::npuw::util::isDQScaleOrZPValue(const std::string& str) {
976+
// Match DynamicQuantize scale/zp parameters for past value cache
977+
static const std::regex pattern(R"(DynamicQuantize/\d+/past_key_values/value/(?:scale|zp))");
978+
return std::regex_match(str, pattern);
979+
}
980+
969981
bool ov::npuw::util::isRestoredPastKeyValueParam(const std::string& str) {
970982
// Match badly handled KVCache states by StatefulToStateless pass for Whisper.
971983
static const std::regex restored_pattern(

‎src/plugins/intel_npu/src/plugin/npuw/util.hpp‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -346,6 +346,11 @@ bool isPastKeyParam(const std::string& str);
346346
// Matches any past value param: contiguous or block-split.
347347
bool isPastValueParam(const std::string& str);
348348

349+
// Detects DynamicQuantize scale/zp parameters for past KV cache.
350+
// Returns true if the parameter name matches the DQ naming pattern.
351+
bool isDQScaleOrZPKey(const std::string& str);
352+
bool isDQScaleOrZPValue(const std::string& str);
353+
349354
// To remove input KV params that got badly matched in StatefulToStateless pass
350355
// in Whisper model.
351356
bool isRestoredPastKeyValueParam(const std::string& str);

‎src/plugins/intel_npu/tests/unit/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -151,6 +151,7 @@ target_sources(${TARGET_NAME} PRIVATE
151151
${CMAKE_CURRENT_SOURCE_DIR}/npuw/resolve_dynamic_quant_storage_types_test.cpp
152152
${CMAKE_CURRENT_SOURCE_DIR}/npuw/orc.cpp
153153
${CMAKE_CURRENT_SOURCE_DIR}/npuw/sdpa_pattern_nodes_test.cpp
154+
${CMAKE_CURRENT_SOURCE_DIR}/npuw/sdpa_pattern_matcher_test.cpp
154155
${CMAKE_CURRENT_SOURCE_DIR}/npuw/subgraph_pipeline.cpp
155156
${CMAKE_CURRENT_SOURCE_DIR}/npuw/subgraph_behavior_infer_test.cpp
156157
${CMAKE_CURRENT_SOURCE_DIR}/npuw/preserve_const_matmul_pattern_test.cpp

0 commit comments

Comments
 (0)