How to replace the attention manchism in the transformconv with flashattn or causal attn?

[HelloWorldLTY]

Hi, I think there should be approaches to optimize the attention calculation in TransformConv codes. How to replace it with multihead attention provided by pytorch with acclecration of flash attn or functional approach like causal attn? Thank a lot.

[HelloWorldLTY]

I tried but still got different output.

[rusty1s]

That's a bit tricky because PyTorch expects dense input tensors of the form [batch_size, num_nodes, num_features], while PyG supports varying number of nodes. The only way I can think of right now to accelerate is via calling MultiheadAttention on a nested tensor (but this doesn't seem to have a backward pass right now, see https://pytorch.org/docs/stable/generated/torch.nn.MultiheadAttention.html).