Skip to content

[NPUW] Support KV cache i8 compression with pyramid attention#36143

Merged
esmirno merged 8 commits into
openvinotoolkit:masterfrom
esmirno:esmirno1/pyramid-attention-kv-i8
Jul 2, 2026
Merged

[NPUW] Support KV cache i8 compression with pyramid attention#36143
esmirno merged 8 commits into
openvinotoolkit:masterfrom
esmirno:esmirno1/pyramid-attention-kv-i8

Conversation

@esmirno

@esmirno esmirno commented May 29, 2026

Copy link
Copy Markdown
Contributor

Details:

  • Add SDPACompressed pattern matcher to isolate i8 dequantized KV cache attention subgraphs (Convert -> Subtract -> Multiply -> Concat)
  • Make Convert ops optional in SDPACompressed to handle models where past KV is already in the expected type
  • Support both asymmetric (with zero-point Subtract) and symmetric (scale-only Multiply) dequantization
  • Propagate DQ scale/zero-point parameters through pyramid attention model creation
  • Add isDQScaleOrZPKey/isDQScaleOrZPValue utility functions for identifying DQ parameters by name
  • Add unit tests for SDPACompressed pattern matching and pyramid attention with i8 KV cache

Tickets:

AI Assistance:

  • AI assistance used: yes
  • Copilot assisted with code implementation, pattern matcher design, and test creation.

@esmirno esmirno requested review from a team as code owners May 29, 2026 19:40
@github-actions github-actions Bot added category: build OpenVINO cmake script / infra category: NPU OpenVINO NPU plugin category: NPUW NPUW plugin labels May 29, 2026
@esmirno esmirno requested a review from intelgaoxiong June 16, 2026 11:55

@intelgaoxiong intelgaoxiong left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

👍LGTM.
Anyway, the block-based KV cache will need to be extended in the future, but I hope to support fp16 KV cache as the first step.

if (dim_iter != past_value_sequence_dims.end()) {
attention._inputs.push_back(ov::npuw::function::Attention::Param{param, dim_iter->second});
}
} else if (ov::npuw::util::isDQScaleOrZPKey(param_name)) {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

A note for the future - instead of names, should we better tag it with rt_info?

if (!past_key_sequence_dims.empty()) {
attention._inputs.push_back(
ov::npuw::function::Attention::Param{param, past_key_sequence_dims.begin()->second});
}

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

what if it is empty? why we would have such a case?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

this case perfectly covers contract for this function API - while for entire flow it is right now not possible combination. I wouldnt call begin() for empty structure anyway, better to have more logs of course here or even controlled exceptions.

std::optional<PyramidAttention> PyramidAttention::from(const std::shared_ptr<ov::Model>& model) {
    // Validate and setup pyramid attention
    auto validation_result = validate_and_setup_pyramid_attention(model); <-- sequence dims gets parsed or died
    if (!validation_result) {
        return std::nullopt;
    }

}
} else if (ov::npuw::util::isDQScaleOrZPValue(param_name)) {
if (!past_value_sequence_dims.empty()) {
attention._inputs.push_back(

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

would these inputs be handled by the pyramid's input ROI (view) logic the same way as past k/v len?

so no other extra handling required?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

I'm expecting so, quant params should match pyramid sequnce len for per-token quantization.

@dmatveev dmatveev added this to the 2026.3 milestone Jun 25, 2026
@esmirno esmirno force-pushed the esmirno1/pyramid-attention-kv-i8 branch 2 times, most recently from 6487d1d to cb91072 Compare June 25, 2026 17:30
@esmirno esmirno requested review from a team as code owners June 25, 2026 17:30
@esmirno esmirno requested review from mlukasze and removed request for a team June 25, 2026 17:30
@github-actions github-actions Bot added category: CPU OpenVINO CPU plugin category: dependency_changes Pull requests that update a dependency file no-match-files and removed category: dependency_changes Pull requests that update a dependency file labels Jun 25, 2026
@esmirno esmirno force-pushed the esmirno1/pyramid-attention-kv-i8 branch from de8e527 to c99494a Compare June 25, 2026 18:30
@github-actions github-actions Bot removed category: CPU OpenVINO CPU plugin no-match-files labels Jun 25, 2026
Comment on lines +324 to +347
isolate_matched(convert1);
isolate_matched(zp_convert1);
isolate_matched(subtract1);
isolate_matched(multiply1);
isolate_matched(concat1);
isolate_matched(unsqueeze1);
isolate_matched(broadcast1);
isolate_matched(reshape1);

isolate_matched(convert2);
isolate_matched(zp_convert2);
isolate_matched(subtract2);
isolate_matched(multiply2);
isolate_matched(concat2);
isolate_matched(unsqueeze2);
isolate_matched(broadcast2);
isolate_matched(reshape2);

isolate_matched(matmul1);
isolate_matched(add);
isolate_matched(softmax);
isolate_matched(matmul2);
isolate_matched(transpose);
isolate_matched(reshape3);

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

If you only need to mark the matched nodes, consider obtaining them using Matcher::get_matched_values()

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

@CuriousPanCake thanks for suggestion, for our need it looks that API cannot be used directly - get_matched_values() usages will require explicit filtering of any_input() nodes, that is useless overcomplication of simple code. So here we have verbose but direct control of what gets isolated versus most extensive filtering / deduplication

// [any_input] → Subtract(zp) → Multiply(scale) → Concat
// instead of the original:
// Convert → Concat
class SDPACompressed : public ov::pass::MatcherPass {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

What's the point of this class?

@esmirno esmirno force-pushed the esmirno1/pyramid-attention-kv-i8 branch 2 times, most recently from a6a2657 to c021878 Compare June 30, 2026 01:12
for (size_t i = 0; i < kc->get_input_size(); ++i) {
const auto& shape = kc->input(i).get_source_output().get_partial_shape();
if (shape.rank().is_static() && shape[kc_axis].is_static()) {
actual_context_length += shape[kc_axis].get_length();

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

I suppose the value of actual_context_length is the same with the original actual_context_length?

@esmirno esmirno added this pull request to the merge queue Jul 2, 2026
Merged via the queue into openvinotoolkit:master with commit cd301a6 Jul 2, 2026
190 of 192 checks passed
@esmirno esmirno deleted the esmirno1/pyramid-attention-kv-i8 branch July 2, 2026 12:17
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

category: build OpenVINO cmake script / infra category: NPU OpenVINO NPU plugin category: NPUW NPUW plugin

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants