Try to unwrap a HuggingFace repeat_kv pattern. HuggingFace's repeat_kv expands KV heads for grouped query attention: hidden_states[:, :, None, :, :].expand(B, n_kv, n_rep, T, D) .clone().reshape(B, n_heads, T, D) In Edge IR this becomes: unsqueez
(cls, node: Node)
| 423 | |
| 424 | @classmethod |
| 425 | def _try_unwrap_repeat_kv(cls, node: Node) -> Optional[Tuple[Node, List[Node]]]: |
| 426 | """Try to unwrap a HuggingFace repeat_kv pattern. |
| 427 | |
| 428 | HuggingFace's repeat_kv expands KV heads for grouped query attention: |
| 429 | hidden_states[:, :, None, :, :].expand(B, n_kv, n_rep, T, D) |
| 430 | .clone().reshape(B, n_heads, T, D) |
| 431 | |
| 432 | In Edge IR this becomes: |
| 433 | unsqueeze_copy(x, 2) → expand_copy → clone → view_copy |
| 434 | |
| 435 | Returns: |
| 436 | (base_node, body_nodes) if pattern matches, else None. |
| 437 | base_node is the original [B, n_kv, T, D] tensor. |
| 438 | body_nodes are the intermediate nodes to absorb. |
| 439 | """ |
| 440 | result = walk_back( |
| 441 | node, |
| 442 | [ |
| 443 | OpStep(op=torch.ops.aten.view.default, nargs=2), |
| 444 | OpStep(op=torch.ops.aten.clone.default, optional=True), |
| 445 | OpStep(op=torch.ops.aten.expand.default, nargs=2), |
| 446 | OpStep(op=torch.ops.aten.unsqueeze.default, nargs=2), |
| 447 | ], |
| 448 | ) |
| 449 | if result is None: |
| 450 | return None |
| 451 | |
| 452 | base, entries = result |
| 453 | _view, _clone, _expand, unsqueeze = entries |
| 454 | |
| 455 | # unsqueeze must be on dim=2 |
| 456 | if unsqueeze.args[1] != 2: |
| 457 | return None |
| 458 | |
| 459 | body = [e for e in entries if e is not None] |
| 460 | return base, body |
| 461 | |
| 462 | @classmethod |
| 463 | def maybe_create(cls, ep: ExportedProgram, head: Node) -> Optional["SDPAHandler"]: |
no test coverage detected