[NPUW] Support KV cache i8 compression with pyramid attention#36143
Conversation
intelgaoxiong
left a comment
There was a problem hiding this comment.
👍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)) { |
There was a problem hiding this comment.
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}); | ||
| } |
There was a problem hiding this comment.
what if it is empty? why we would have such a case?
There was a problem hiding this comment.
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( |
There was a problem hiding this comment.
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?
There was a problem hiding this comment.
I'm expecting so, quant params should match pyramid sequnce len for per-token quantization.
6487d1d to
cb91072
Compare
de8e527 to
c99494a
Compare
| 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); |
There was a problem hiding this comment.
If you only need to mark the matched nodes, consider obtaining them using Matcher::get_matched_values()
There was a problem hiding this comment.
@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 { |
There was a problem hiding this comment.
What's the point of this class?
a6a2657 to
c021878
Compare
| 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(); |
There was a problem hiding this comment.
I suppose the value of actual_context_length is the same with the original actual_context_length?
cd301a6
Details:
Tickets:
AI Assistance: