[CPU] Allow Gather on KV cache paths in MatchSdpaKvCache#35162
[CPU] Allow Gather on KV cache paths in MatchSdpaKvCache#35162CuriousPanCake wants to merge 1 commit intoopenvinotoolkit:masterfrom
Conversation
There was a problem hiding this comment.
Pull request overview
Fixes a CPU plugin graph-optimizer pattern match so MatchSdpaKvCache still runs when a Gather node appears as an additional consumer on KV-cache paths after StatefulSDPAFusion, preventing missing KV state assignment and a null-state crash during inference.
Changes:
- Extends
MatchSdpaKvCache’s eligibility check to allowType::GatheramongMemoryInput’s children (in addition toScaledDotProductAttentionandShapeOf).
| for (auto&& item : childEdges) { | ||
| auto childNode = item->getChild(); | ||
| if (none_of(childNode->getType(), Type::ScaledDotProductAttention, Type::ShapeOf)) { | ||
| if (none_of(childNode->getType(), Type::ScaledDotProductAttention, Type::ShapeOf, Type::Gather)) { |
There was a problem hiding this comment.
[HIGH] This change fixes a crash-prone behavioral bug in KV-cache state assignment, but the PR doesn’t add a regression test covering the newly accepted pattern (MemoryInput having an additional Gather consumer alongside SDPA/ShapeOf). Please add a CPU plugin unit/functional test that builds a model where ReadValue/MemoryInput feeds the fused SDPA-with-KV-cache path and also has a sibling Gather consumer, and asserts compilation + inference works (no null-state crash) and the KV states are assigned/usable.
|
This PR is a WA. The rootcause was additionally investigated and a solution merged within bcbe69a |
Details:
MatchSdpaKvCache is not executed due to the presence of Gather among the Concat's children on KV cache paths after StatefulSDPAFusion. This prevents
MemoryInputSDPAcreation, leaving KV states unassigned and causinga null state crash at inference.
Tickets:
AI Assistance:
Signed-off-by: Andrii Staikov andrii.staikov@intel.com