Pattern for Scaled Dot Product Attention with optional GQA. Matches: scaled_dot_product_attention Optionally with repeat_interleave for grouped query attention.
| 391 | |
| 392 | @REGISTRY.register_pattern(name="SDPA") |
| 393 | class SDPAHandler(PatternHandler): |
| 394 | """ |
| 395 | Pattern for Scaled Dot Product Attention with optional GQA. |
| 396 | |
| 397 | Matches: scaled_dot_product_attention |
| 398 | Optionally with repeat_interleave for grouped query attention. |
| 399 | """ |
| 400 | |
| 401 | def __init__( |
| 402 | self, |
| 403 | head: Node, |
| 404 | body: List[Node], |
| 405 | q_node: Node, |
| 406 | k_node: Node, |
| 407 | v_node: Node, |
| 408 | ): |
| 409 | super().__init__(head, body) |
| 410 | self.q_node = q_node |
| 411 | self.k_node = k_node |
| 412 | self.v_node = v_node |
| 413 | |
| 414 | @classmethod |
| 415 | def _parse_sdpa_args_and_kwargs(cls, sdpa_node: Node): |
| 416 | q, k, v = sdpa_node.args[0:3] |
| 417 | attn_mask = sdpa_node.args[3] if len(sdpa_node.args) > 3 else None |
| 418 | dropout_p = sdpa_node.args[4] if len(sdpa_node.args) > 4 else 0.0 |
| 419 | is_causal = sdpa_node.args[5] if len(sdpa_node.args) > 5 else False |
| 420 | enable_gqa = sdpa_node.args[6] if len(sdpa_node.args) > 6 else False |
| 421 | scale = sdpa_node.kwargs.get("scale", None) |
| 422 | return q, k, v, attn_mask, dropout_p, is_causal, scale, enable_gqa |
| 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 |