Utilizing tensor cores for `mma!`

[AntonOresten]

I have not worked with e.g. CUDA.WMMA before, but perhaps as part of the CUDA extension, a method can be defined for leveraging the tensor cores instead of the mma! which I assume only uses e.g. CUDA cores.
I did not see any mention of MMA/WMMA in KA.jl.

[AntonOresten]

Hm... mma! is thread-level though at the moment, isn't it...
EDIT: I guess that's not a problem though.

[AntonOresten]

It seems CUDA.WMMA may not support BFloat16

[AntonOresten]

I got a AB+C kernel working using WMMA, 8x slower than cuBLAS, but that was all grabbing straight from global memory instead of shmem.

[AntonOresten]

I was able to use a wmma! function to make a flash FFN forward pass ~25% faster. Something like cuTile might be useful here though. See JuliaGPU/CUDA.jl#2991

[AntonOresten]

#26 adds support through CUDA.WMMA

[AntonOresten]

There is now https://github.com/JuliaGPU/cuTile.jl (unregistered), but it seems it can't be used from within a regular CUDA kernel. Possibly something like https://cuda.juliagpu.org/stable/lib/driver/#CUDA.CuGraph could be used.