Questions About Example 35: Gemm + Softmax Fusion

[FabianSchuetze]

Thanks for this wonderful repo! It's a pleasure to work with it and learn from it.

I have a few questions regarding example 35 (gemm + softmax fusion):

• Why is there device memory allocated for block_Sum and block_Norm, and to a certain extent also to block_D? I expected these variables to be stored in shared memory to avoid the memory traffic with main memory.
• Why is there a need for a separate SoftmaxApplyKernel and an ApplyFinalReduction kernel? I would have expected that the EpilgoueVisitorSoftmax is enough.
• Finally, can the custom GemmWithEpilogueVistor be replaced by the general GemmWithEpilogueVistor? I tried to replace them, and the compilation failed. I tried to address the compiler issues, but didn't fully succeed.

[FabianSchuetze]

I've worked more on this and I think I can answer myself. For those who care:

1. I was wrong, the allocations are indeed necessary. The main memory buffers are used for commutation between the threadblocks.
2. The name EpilgoueVisitorSoftmax is a bit misleading in my opinion. This visitor does compute the max element and sum of the elements per row in threadblock. No normalization is done and other kernels are needed to compute the softmax. As the softmax is a form of reduction i wonder if the other reduction-epilogues also operate on a per threadblock basis and not across entire D.
3. I tried various combinations but using any Gemm Kernel from include fails for various reasons and the custom kernel needs to be used.