[NPUW] Support multiple outputs in LLMCompiledModel

[AsyaPronina]

Details:

• PR aims to support OmniThinker model.
• Fixed partitioning issue with cut LM Head for OmniThinker. If LM head cutting is ON, then output embeddings (not logits) became a Result from the prefill/kvcache model. However, last operation before LM head has had already a connected Result node corresponding to the second output of the original OmniThinker model. Thus one output layer has two Result nodes what didn't work with partitioning (but worked without partitioning).
• Improved CutLMHead transformation.
  • Addes check that Result's output name is "logits". This check is sufficiently reliable check for finding exactly logits output, because:
    • LLMInferRequest always rely on "logits" name to get logits from prefill/kvcache models.
    • Following Exporter configs: OnnxConfig, OnnxConfigWithPast, TextDecoderOnnxConfig and TextDecoderWithPositionIdsOnnxConfig from optimum-onnx name LLM output with "logits": https://github.com/huggingface/optimum-onnx/blob/main/optimum/exporters/onnx/base.py#L137, https://github.com/huggingface/optimum-onnx/blob/main/optimum/exporters/onnx/config.py#L100
    • Most of optimum-intel OpenVINO Exporter configs are derived from the configs above: https://github.com/huggingface/optimum-intel/blob/main/optimum/exporters/openvino/model_configs.py#L684
    • optimum-intel export() function set names for output tensors from Exporter config: https://github.com/huggingface/optimum-intel/blob/main/optimum/exporters/openvino/convert.py#L442-L445

Tickets:

• ticket-id

[dmatveev]

If LM head cutting is ON, then output embeddings (not logits) became a Result from the prefill/kvcache model. However, last operation before LM head has had already a connected Result node corresponding to the second output of the original OmniThinker model

If there's already a Result, why don't you use one in your 3mp?

[dmatveev]

I believe this change fixes the issue when a single Node has two Results connected to it.. And this happens when you introduce one more Result to feed the 3rd model. Can this be avoided?

If our partitioning doesn't work for that case, I'd probably fix it later

[AsyaPronina]

I understand the concern, however, CutLMHead transformations happens before the model is split into prefill and generate. When we work with the prefill model, we also add a Slice but before only one Result node. So two Result nodes got separated: one is now after the Slice , another one is still connected to the layer before LM head. However, we don't need Slice for generate model as it outputs 1 token already. And only here we face the issue with two Results from one layer before the LM head.
If to merge two Results into one in advance, then Slice will be added to the both results..

[AsyaPronina]

After testing it turns out that preserving changes in partitioning is a more safe approach. Merging multiple Result nodes per output layer for generate model works, however this one Result node would have multiple names (from multiple Result-s node to catch their meanings). But, in LLMInferRequest, on the contrary, we are using output_port->get_any_name() to get any tensor name but only one for mapping of names to outputs. This may cause issues. So, preserving partitioning changes currently.

[GuoliangShiIntel]

@AsyaPronina I have tested this change on both the Eagle3 and Omni pipelines, and both work well.

I'd like to leave an additional comment here in case it's related to the "multi-outputs" feature. Since "multi-outputs" may result in incorrect KV cache redirection, I made some related changes in the previous Eagle3 PR.

I noticed that Eugene added a comment after the PR was merged. I'm not sure if there's anything we should address regarding this?

[GuoliangShiIntel]

@AsyaPronina And another question, is this PR plan for OV 26.0 release? if yes, i will also change for eagle3 pipeline to align this PR.

[dmatveev]

@GuoliangShiIntel there's a lot PR checks failing here for some reason..

[AlexanderKalistratov]

Linux builds failed here:

error: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::shared_ptr<ov::Node> >::size_type' {aka 'long unsigned int'} [-Werror=sign-compare]
for (auto i = 1; i < maybe_results.size(); ++i) {

[AsyaPronina]

Thanks a lot!

[AsyaPronina]

@AsyaPronina I have tested this change on both the Eagle3 and Omni pipelines, and both work well.

I'd like to leave an additional comment here in case it's related to the "multi-outputs" feature. Since "multi-outputs" may result in incorrect KV cache redirection, I made some related changes in the previous Eagle3 PR.

I noticed that Eugene added a comment after the PR was merged. I'm not sure if there's anything we should address regarding this?

Hello @GuoliangShiIntel!
Thanks a lot for thorough review and testing of pipelines!

1. Yes, comment about KVCache redirection and Eugene's comment about empty KV inputs removal are valid.
   • Regarding KVCache redirection: for now your solution for Eagle3 additional outputs will work perfect. As long as we are only redirecting outputs with present_key or present_value names, issues shouldn't arise.
   • Regarding KVCache redirection in @esmirno Eugene's PR: Eugene refined KVCache redirection logic to use pattern matching. These patterns still won't match neither additional outputs of Eagle3 nor of Omni.
   • Regading empty KV inputs removal: Removal of empty KV inputs is based on pattern matching. This pattern matching seeks for param->concat pattern which shouldn't match neither Eagle3 output (add, add, add->concat) nor OmniThinker (no concat).

[GuoliangShiIntel]

LGTM