Class ov::pass::AttentionMaskShapeReplacer

class AttentionMaskShapeReplacer : public ov::pass::MatcherPass

Replaces ShapeOf(attention_mask) with ShapeOf(input_source) in the broadcast-shape branch that feeds rotary MatMul.

Some models build a Broadcast target shape from attention_mask and then use that Broadcast as the first input of MatMul, while the second input is derived from position_ids through an optional Unsqueeze/Convert chain. In PagedAttention mode, attention_mask is removed, therefore this shape branch must be rewired to an always-present source (input_ids or inputs_embeds).

Matched topology:

┌─────────────────┐ │ attention_mask │ │ (Parameter) │ └────────┬────────┘ │ ┌────┴────┐ ┌──────────────┐ │ ShapeOf │ │ position_ids │ └────┬────┘ │ (Parameter) │ │ └──────┬───────┘ ┌────┴────┐ ┌─────┴──────┐ │ Gather │ │ Unsqueeze │ (optional) └────┬────┘ └─────┬──────┘ │ ┌────┴────┐ ┌────┴────┐ │ Convert │ (optional) │ Concat │ └────┬────┘ └────┬────┘ │ │ │ ┌────┴────┐ │ │Broadcast├───────────────┘ └────┬────┘ │ ┌────┴────┐ │ MatMul │ └─────────┘

Rewrite:

Replace ShapeOf(attention_mask) with ShapeOf(input_source), where input_source is input_ids or inputs_embeds.

After:

┌─────────────────┐ │ input_source │ │ (Parameter) │ └────────┬────────┘ │ ┌────┴────┐ ┌──────────────┐ │ ShapeOf │ │ position_ids │ └────┬────┘ │ (Parameter) │ │ └──────┬───────┘ ┌────┴────┐ ┌─────┴──────┐ │ Gather │ │ Unsqueeze │ (optional) └────┬────┘ └─────┬──────┘ │ ┌────┴────┐ ┌────┴────┐ │ Convert │ (optional) │ Concat │ └────┬────┘ └────┬────┘ │ │ │ ┌────┴────┐ │ │Broadcast├───────────────┘ └────┬────┘ │ ┌────┴────┐ │ MatMul │ └─────────┘

The rewrite is applied only when all of the following hold: 1) Gather axis is 0 (or -1 for a 1D ShapeOf output). 2) Gather indices are exactly {0}. 3) input_source rank is static and at least 1.

These checks guarantee that the selected dimension is equivalent after rewiring and avoid altering unsupported shape semantics.