Distributed Neighborhood Attention S2

[azrael417]

This MR adds distributed Neighborhood Attention S2 support and fixes some issues in the existing attention kernel.

• existing serial attention kernel preallocated an output tensor which was too large when using attention based downsampling. This is fixed
• existing serial attention kernel does not produce the correct v gradient when used in upsampling. We will fix that next
• this MR adds distributed neighborhood attention along with some new tests for the feature. This kernel does not support up- or downsampling yet