Repository: microsoft/onnxruntime
File: onnxruntime/python/tools/transformers/fusion_bart_attention.py (Python)

Code before:
        if (
            encoder_attention
            or decoder_self_attention
            or decoder_cross_attention
            or decoder_self_attention_with_past
            or decoder_cross_attention_with_past
        ):
            attention_last_node = reshape_qkv
            num_heads, hidden_size = self.get_num_heads_and_hidden_size(reshape_q)

            # Fall back to user-specified values when detected values are invalid
            # (e.g., SDPA models use -1 in reshape shapes for dynamic dimensions).
            if (num_heads <= 0 or hidden_size <= 0) and self.num_heads > 0 and self.hidden_size > 0:
                num_heads = self.num_heads
                hidden_size = self.hidden_size

            if num_heads <= 0 or hidden_size <= 0 or (hidden_size % num_heads) != 0:
                logger.debug("fuse_attention: failed to detect num_heads or hidden_size")
                return

            new_node = None
            if decoder_self_attention_with_past or decoder_cross_attention or decoder_cross_attention_with_past:
                # Note: Decoder attention with past key and past value is fused as multi-head attention
                # rather than attention because multi-head attention supports separate past key and past
                # value whereas attention supports concatenated past key and past value.
                new_node = (
                    self.create_multihead_attention_node(
                        q_matmul=matmul_q,
                        k_matmul=matmul_k if decoder_cross_attention or decoder_self_attention_with_past else past_k,
                        v_matmul=matmul_v if decoder_cross_attention or decoder_self_attention_with_past else past_v,
                        q_add=add_q,
                        k_add=add_k if decoder_cross_attention or decoder_self_attention_with_past else None,
                        v_add=add_v if decoder_cross_attention or decoder_self_attention_with_past else None,
                        num_heads=num_heads,
                        hidden_size=hidden_size,
                        output=attention_last_node.output[0],
                        unidirectional=causal_mask,
                        past_k=past_k if decoder_self_attention_with_past else "",
                        past_v=past_v if decoder_self_attention_with_past else "",
                        present_k=present_k,
                        present_v=present_v,
                    )
                    if self.use_multi_head_attention
                    else None
                )
            else:
                # Temporarily set multi-head attention flag to false
                use_multi_head_attention_ground_truth = self.use_multi_head_attention
                self.use_multi_head_attention = False
                new_node = self.create_attention_node(
                    mask_index=None,

Review comment (bug):
When num_heads/hidden_size fallback is used, there's no logging to indicate this happened. Consider adding a debug log.

Code after:
                return
            assert len(mask_nodes) > 0

        if (
            encoder_attention
            or decoder_self_attention
            or decoder_cross_attention
            or decoder_self_attention_with_past
            or decoder_cross_attention_with_past
        ):
            attention_last_node = reshape_qkv
            num_heads, hidden_size = self.get_num_heads_and_hidden_size(reshape_q)

            # Fall back to user-specified values when detected values are invalid
            # (e.g., SDPA models use -1 in reshape shapes for dynamic dimensions).
            if (num_heads <= 0 or hidden_size <= 0) and self.num_heads > 0 and self.hidden_size > 0:
                logger.debug(
                    "fuse_attention: reshape dims invalid (num_heads=%d, hidden_size=%d), "
                    "falling back to user-specified num_heads=%d, hidden_size=%d",
                    num_heads,
                    hidden_size,
                    self.num_heads,
                    self.hidden_size,
                )
                num_heads = self.num_heads
                hidden_size = self.hidden_size

            if num_heads <= 0 or hidden_size <= 0 or (hidden_size % num_heads) != 0:
                logger.debug("fuse_attention: failed to detect num_heads or hidden_size")
                return

            new_node = None
            if decoder_self_attention_with_past or decoder_cross_attention or decoder_cross_attention_with_past:
                # Note: Decoder attention with past key and past value is fused as multi-head attention
                # rather than attention because multi-head attention supports separate past key and past
                # value whereas attention supports concatenated past key and past value.
                new_node = (
                    self.create_multihead_attention_node(
                        q_matmul=matmul_q,
                        k_matmul=matmul_k if decoder_cross_attention or decoder_self_attention_with_past else past_k,
                        v_matmul=matmul_v if decoder_cross_attention or decoder_self_attention_with_past else past_v,
                        q_add=add_q,
                        k_add=add_k if decoder_cross_attention or decoder_self_attention_with_past else None,
                        v_add=add_v if decoder_cross_attention or decoder_self_attention_with_past else None,
                        num_heads=num_heads,
                        hidden_size=hidden_size,
                        output=attention_last_node.output[0],
                        unidirectional=causal_mask,
                        past_k=past_k if decoder_self_attention_with_past else "",
                        past_v=past_v if decoder_self_attention_with_past else "",
                        present_k=present_k,