Is there any doc about fmha

NVIDIA / apex

A PyTorch Extension: Tools for easy mixed precision and distributed training in Pytorch

BSD 3-Clause "New" or "Revised" License

8.34k stars 1.39k forks source link

Open wukong1992 opened 1 year ago

wukong1992 commented 1 year ago

Is there any doc about fmha implementation. Read src, so many confusing [code:] (https://github.com/NVIDIA/apex/blob/6309120bf4158e552884427b18ebafc657d418df/apex/contrib/csrc/fmha/src/fmha_fprop_kernel_1xN.h#L515) Why thers is griddim.x not blockdim.x. It's so bad experience to read the code.