Implement AT_DISPATCH for CUDA attention kernel dispatch

[azrael417]

We should use the torch dispatchers to support auto casting in the future. Everything works currently because the C++ wrappers do upcasting, but the cleaner solution would be to use the dispatchers.

[azrael417]

this is done