[stable-diffusion v2] ONNX pattern matching for mhca and mha plugins is not working

Description

The pattern matching used in the diffusion demo for replacing memory head attention with the optimized plugins is not working with stable diffusion v2 in the demo-diffusion.

Resulting to this TensorRT error :

[E] 10: [optimizer.cpp::computeCosts::3712] Error Code 10: Internal Error (Could not find any implementation for node {ForeignNode[MatMul_450...MatMul_457]}.)
[E] 2: [builder.cpp::buildSerializedNetwork::738] Error Code 2: Internal Error (Assertion engine != nullptr failed. )
[!] Invalid Engine. Please ensure the engine was built correctly.

I would like to know if it is something easily patchable on my side in the function mha_mhca_detected from demo-diffusion

Environment

TensorRT branch release/8.5
Docker container image : nvcr.io/nvidia/tensorrt:22.10-py3

Relevant Files

Files models.py and demo_diffusion.py with stable-diffusion-2 support are in the files.zip archive
After investigation, the graphs are not identical. For example, the projection layers in v1 use Conv instead of Linear in v2.

With fhma plugin support deactivated, we can observe that the ONNX graphs are not the same between the two version.
An additional MatMul operation is added as input for the attention part. I suspect this additional operator being responsible of the issue.

Subgraph-v1, Attention Input:

Subgraph-v2, Attention Input:

Also, I am not sure where this part should applied in the graph :

    def mha_mhca_detected(self, node, mha):
        # Go from V GEMM down to the S*V MatMul and all way up to K GEMM
        # If we are looking for MHCA inputs of two matmuls (K and V) must be equal.
        # If we are looking for MHA inputs (K and V) must be not equal.
        if node.op == "MatMul" and len(node.outputs) == 1 and \
            ((mha and len(node.inputs[0].inputs) > 0  and node.i().op == "Add") or \
            (not mha and len(node.inputs[0].inputs) == 0)):

            dynamic_batch = (node.o().op == 'Shape')
            o = node.o(1) if dynamic_batch else node.o()
            if o.op == "Reshape" and \
                o.o().op == "Transpose" and \
                o.o().o().op == "Reshape" and \
                o.o().o().o().op == "MatMul" and \
                o.o().o().o().i(0).op == "Softmax" and \
                o.o().o().o().i(1).op == "Reshape" and \
                o.o().o().o().i(0).i().op == "Mul" and \
                o.o().o().o().i(0).i().i().op == "MatMul" and \
                o.o().o().o().i(0).i().i().i(0).op == "Reshape" and \
                o.o().o().o().i(0).i().i().i(1).op == "Transpose" and \
                o.o().o().o().i(0).i().i().i(1).i().op == "Reshape" and \
                o.o().o().o().i(0).i().i().i(1).i().i().op == "Transpose" and \
                o.o().o().o().i(0).i().i().i(1).i().i().i().op == "Reshape" and \
                o.o().o().o().i(0).i().i().i(1).i().i().i().i().op == "MatMul" and \
                node.name != o.o().o().o().i(0).i().i().i(1).i().i().i().i().name:
                # "len(node.outputs) == 1" to make sure we are not in the already fused node
                node_q = o.o().o().o().i(0).i().i().i(0).i().i().i()
                node_k = o.o().o().o().i(0).i().i().i(1).i().i().i().i()
                node_v = node
                final_tranpose = o.o().o().o().o(1).o() if dynamic_batch else o.o().o().o().o().o()
                # Sanity check to make sure that the graph looks like expected
                if node_q.op == "MatMul" and final_tranpose.op == "Transpose":
                    return True, dynamic_batch, node_q, node_k, node_v, final_tranpose
        return False, False, None, None, None, None

If you have additional information to explain this pattern matching, it could be great.

Steps To Reproduce

pip3 install -r demo/Diffusion/requirements.txt
Follow the steps from hugging face : https://huggingface.co/stabil
ityai/stable-diffusion-2
pip3 install --upgrade git+https://github.com/huggingface/diffusers.git transformers accelerate scipy
Run the demo:
LD_PRELOAD=${PLUGIN_LIBS} python3 demo/Diffusion/demo-diffusion.py "" -v

files.zip