rfc: graph: propose to support Grouped-query Attention

[gyhintel]

Description

This is to propose to support Grouped-query Attention in oneDNN Graph API.
Link to the rendered document.

[dzarukin]

@gyhintel, thanks for the RFC. Some questions:

1. If the user specifies a graph with different dimensions for K x Q, H dim versus G dim as in the example, do you see any issues with expanding pattern matcher / underneath implementation to perform an extra check for H dividing G, and if yes, perform the logic as described in Option 2 in the document?
2. If (1) can be done, would it mean that the user must modify their Graph, or dnnl::graph, to match our pattern for GQA?
3. If (2) forces the user to do modification, will adding API from Option 1 in the document help to expand out pattern matcher so that there're no extra actions needed from the user side?

[gyhintel]

1. If the user specifies a graph with different dimensions for K x Q, H dim versus G dim as in the example, do you see any issues with expanding pattern matcher / underneath implementation to perform an extra check for H dividing G, and if yes, perform the logic as described in Option 2 in the document?

1. If there is no Reshape above the typical SDPA, Query(in shape (N, H, S, D)) and Key/Value (in shape (N, G, S, D)) have different head number dimension and cannot perform dot-product directly. This is the doc: "For example src can be broadcasted to wei, if the corresponding dimension in src is 1 (and vice versa). ".
2. We can extend the MatMul broadcasting rules to support group broadcast. This is option 3. In this situation, we need to perform an extra check for H dividing G. There should be no issues.

2. If (1) can be done, would it mean that the user must modify their Graph, or dnnl::graph, to match our pattern for GQA?

Yes, it means that the pattern cannot be used to optimize a framework graph directly. Users will have to map their GQA implementation graph to our pattern. This is the second cons of option 2.

3. If (2) forces the user to do modification, will adding API from Option 1 in the document help to expand out pattern matcher so that there're no extra actions needed from the user side?

In the current Pytorch implementation, there are no extra actions from their side. But if the implementation in the community changes, still needs to handle the new implementation. This is the second cons of option1.

[petercad]

If this turns out to be a serious con, it would be reasonable to add a pass to match the Option 1 subgraph and convert it to the Option 2 subgraph, right?

[gyhintel]

1. If it is a serious con, we need to implement option 1 adding new ops and new patterns. It is a backend implementation that matches the Option 1 subgraph and converts it to the Option 2 subgraph. We can also implement it in other ways in the backend.
2. If the pass can be done on the framework side, we only need to implement option 2.

[TaoLv]

We will have to support and match the subgraph in Option 1 once the request pops up. With that, oneDNN will support and maintain several different patterns for the same GQA functionality. Maybe it's not an issue as even for now we choose to Option 1 as the initial step, the pattern may still change in the future as mentioned in the cons of Option 1.

[TaoLv]

I'd like to add the link here for reference: https://docs.nvidia.com/deeplearning/cudnn/latest/api/cudnn-graph-library.html#cudnn-backend-operation-matmul-descriptor .

[gyhintel]

Added, thanks.

[TaoLv]

Another con here maybe that we rely on oneDNN matmul primitive kernels for reference implementation and testing in benchdnn which do not support the new broadcasting rule. Extending the broadcast semantics on graph side will also request additional effort for reference implementation and testing.

[gyhintel]

Added, thanks.

[TaoLv]

FYI - the PyTorch PR just got merged this week: pytorch/pytorch#132689

[dzarukin]

Any clarity at this point if users are fine with implementing Option 2 on their side, or Option 1 must be implemented instead?

[gyhintel]

@chunyuan-w, @sanchitintel, Could you help take a look at this RFC? thanks!

[ElaineBao]

Suggested change

	| B1 x 1 x B3 x M x K | B1 x B2 x 1 x M x K | B1 x B2 x B3 x M x N |
	| B1 x 1 x B3 x M x K | B1 x B2 x 1 x K x N | B1 x B2 x B3 x M x N |

[gyhintel]

Fixed, thanks.