[PyTorch] Reduce the CPU overheads of GroupedLinear

[yaox12]

Description

Try to reduce the CPU overheads of GroupedLinear by:

1. Using fused_multi_cast_transpose instead of iterating cast_transpose_fused in a for-loop.
   a. Changed the API of fused_multi_cast_transpose to avoid index-select ops in PyTorch.
   b. Allocate output tensors in CPP.
2. Use at::cuda::current_device(), which has a cache, to get current device id to avoid cudaGetDriverEntryPoint calls.
3. Reduce torch.Tensor() calls.
4. TODO.

Fix:

1. Zero grad_bias in fused_cast_transpose_bgrad when input is empty.
2. Typos.

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refractor

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[yaox12]

@timmoon10 Can you help review this PR?

[timmoon10]

/te-ci pytorch

[timmoon10]

Overall LGTM. What kinds of speedups are you seeing? It would be helpful to see what optimizations had the biggest impact so we can apply it to other PyTorch modules.

[yaox12]

Overall LGTM. What kinds of speedups are you seeing? It would be helpful to see what optimizations had the biggest impact so we can apply it to other PyTorch modules.

Here're some numbers for num_gemms=6 and small gemm sizes.

	fwd (us)	fwd speedup
baseline	750.075	/
+fused_multi_cast_transpose	662.756	1.13
+cpp alloc	615.818	1.22
+Reduce cudaGetDriverEntryPoint and torch.Tensor()	604.722	1.24

[timmoon10]

/te-ci pytorch