Build and codegen for `Back2Back GEMM` dataflow.

TiledTensor / ThrillerFlow

ThrillerFlow is a Dataflow Analysis and Codegen Framework written in Rust.

MIT License

5 stars 1 forks source link

Build and codegen for `Back2Back GEMM` dataflow. #6

Open KuangjuX opened 1 month ago

KuangjuX commented 1 month ago

Back2Back GEMM is an important kernel, and it is the core of flash attention, so it is necessary to analyze its dataflow and generate it with the help of the dataflow.

KuangjuX commented 1 month ago

Back2Back GEMM is similar to GEMM, but at the register level, it first performs matrix multiplication on the two input matrices A and B, and then performs matrix multiplication on the third matrix. It's worth noting that the mapping of matrices A, B, and C is different. For matrices A and B, the k dimension needs to be split over time, while for matrix C, the p dimension is mapped to the thread block, resulting in a different nested loop structure.