Class ov::pass::EliminateDropBatch#
-
class EliminateDropBatch : public ov::pass::MatcherPass#
After SDPAToPagedAttention flattens position_ids, the batch dimension no longer exists, so an aten::select(dim=0, index=0) that dropped it (lowered to Gather(index=0, axis=0)) becomes redundant.
This transformation detects Parameter(name == “position_ids”) -> Unsqueeze(optional) -> Convert(optional) -> Gather(index=0, axis=0) and replaces the Gather with a Reshape(-1), which flattens whatever shape is produced by the optional Unsqueeze/Convert back to a 1D tensor. Since position_ids no longer carries a real batch dimension to select from, this is equivalent to the original select but no longer depends on a batch axis being present.
We change from this: to this:
┌──────────────┐ ┌──────────────┐ │ position_ids │ │ position_ids │ └──────┬───────┘ └──────┬───────┘ │ │ ┌──────┴──────┐ ┌──────┴──────┐ │ Unsqueeze │ (optional) │ Unsqueeze │ (optional) └──────┬──────┘ └──────┬──────┘ │ │ ┌──────┴──────┐ ┌──────┴──────┐ │ Convert │ (optional) │ Convert │ (optional) └──────┬──────┘ └──────┬──────┘ │ │ ┌──────┴──────┐ ┌──────┴──────┐ │Gather(idx=0,│ │ Reshape │ │ axis=0) │ │ (-1) │ └─────────────┘ └─────────────┘