Skip to content

[HybridAttention] Support past key values in attention_mask for stateful inference#1641

Open
apaniukov wants to merge 9 commits intohuggingface:mainfrom
apaniukov:fix-granite-moe-hybrid-state
Open

[HybridAttention] Support past key values in attention_mask for stateful inference#1641
apaniukov wants to merge 9 commits intohuggingface:mainfrom
apaniukov:fix-granite-moe-hybrid-state

Conversation

@apaniukov
Copy link
Copy Markdown
Contributor

@apaniukov apaniukov commented Mar 19, 2026

What does this PR do?

GraniteMoeHybrid crashes with a shape mismatch error:

RuntimeError: Exception from src/inference/src/cpp/infer_request.cpp:224:
Exception from src/plugins/intel_cpu/src/node.cpp:787:
[CPU] Add node with name '__module.model/aten::add/Add_1' Check 'input_shape[j] == 1' failed at src/plugins/intel_cpu/src/shape_inference/custom/eltwise.cpp:52:
Eltwise shape infer input shapes dim index: 3 mismatch

The RuntimeError at the Add node that combines the causal mask with the position-expanded attention_mask:

  • port 0 (causal mask): [batch, 1, seq_len, seq_len]
  • port 1 (attn mask): [batch, 1, 1, total_seq_len]
    where seq_len != total_seq_len during 2nd+ generation steps when seq_len > 1 (e.g. second prefill stage).

The issue caused by patched _update_causal_mask method. Hybrid models that doesn't patch it (Qwen3-Next and LFM2) work fine.

Fix:

  • Update inputs property to use "past_sequence_length + sequence_length" for attention_mask for -with-past conversion
  • [unrelated] remove unnecessary triu operation

Fixes # (issue)

Before submitting

  • This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
  • Did you make sure to update the documentation with your changes?
  • Did you write any new necessary tests?

@apaniukov apaniukov changed the title Support past key values in attention_mask for stateful inference [GraniteMoeHybrid] Support past key values in attention_mask for stateful inference Mar 19, 2026
@apaniukov apaniukov changed the title [GraniteMoeHybrid] Support past key values in attention_mask for stateful inference [HybridAttention] Support past key values in attention_mask for stateful inference Mar 19, 2026
@HuggingFaceDocBuilderDev
Copy link
Copy Markdown

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Copy link
Copy Markdown
Collaborator

@rkazants rkazants left a comment

Choose a reason for hiding this comment

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

can you please add test to reproduce this case?

@apaniukov
Copy link
Copy Markdown
Contributor Author

can you please add test to reproduce this case?

Done.

@apaniukov apaniukov requested a review from rkazants March 19, 2026 18:30
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants