Grouped swizzle kernel

[ptrendx]

Develop a grouped swizzle kernel for the GroupedTensor type to support efficient memory layout transformations in MoE training. Ensure compatibility with device-supplied problem sizes.

[vthumbe1503]

PR merged which adds support for Grouped Tensor swizzling with uniform input sizes(weights).

Still needs to be done for variable input sizes.

[vthumbe1503]

#2769 also adds additional cleanups and optimizations. Optimization involves swizzling only along the direction(rowise/colwise) where the GEMM needs it.

Can we close this @ptrendx or do we need the support for activation(with variable first_dims) as well?

Since we have the group_quantize+swizzle fusion, users can generally swizzle their activations directly using that. Do we have a use-case for activation/grad_output swizzling?

[vthumbe1503]

Synced with Przemek offline. This issue can be closed