add GatherSliceToSplitFusion and Unittest

[rui-ren]

Multi Query Attention Optimization

in multi-query attention

batch_size, seq_length, three_times_hidden_size = fused_qkv.shape
fused_qkv = fused_qkv.view(batch_size, seq_length, self.num_heads + 2, self.head_dim)
return fused_qkv[..., :-2, :], fused_qkv[..., [-2], :], fused_qkv[..., [-1], :]

which can be optimized to

batch_size, seq_length, three_times_hidden_size = fused_qkv.shape
fused_qkv = fused_qkv.view(batch_size, seq_length, self.num_heads + 2, self.head_dim)
(query, key, value) = fused_qkv.split([self.num_heads, 1, 1], dim=2)
return query, key, value

this optimization can be validated from nsight profiling and perf benchmarking.

As such, This PR is to Optimize the Gather/Gather/Slice Ops to Split Kernel.

Optimization Target

As 2 Gather and 1 Slice Kernels are time consuming for backward prop, it would be efficient to use 1 Split Kernel

Example

• Before Fusion

• After Fusion

Perf Gain

After the optimization, there will have ~7% perf gain.

The Transpose Kernel can be fused too, will update it in next PR. However, after testing Transponse Ops fusion on Falcon model, there is no perf gain. Will not create a new PR.

[askhade]

LGTM