A fused apply_rotary_pos_emb implementation for Megatron-Core

[yaox12]

This is a fused apply_rotary_pos_emb implementation for Megatron-Core.

In my preliminary benchmark, it gives 2x - 4x speedup over the unfused version. batch_size=2 and head_num=64 are fixed.

dtype=torch.float32, seq_length=2048, hidden_size=128, rotary_percent=0.5
unfused rope: 0.45 ms
fused rope: 0.14 ms

dtype=torch.float32, seq_length=2048, hidden_size=128, rotary_percent=1.0
unfused rope: 0.67 ms
fused rope: 0.15 ms

dtype=torch.float32, seq_length=2048, hidden_size=256, rotary_percent=0.5
unfused rope: 0.84 ms
fused rope: 0.27 ms

dtype=torch.float32, seq_length=2048, hidden_size=256, rotary_percent=1.0
unfused rope: 1.3 ms
fused rope: 0.3 ms

dtype=torch.float32, seq_length=4096, hidden_size=128, rotary_percent=0.5
unfused rope: 0.85 ms
fused rope: 0.23 ms

dtype=torch.float32, seq_length=4096, hidden_size=128, rotary_percent=1.0
unfused rope: 1.3 ms
fused rope: 0.3 ms

dtype=torch.float32, seq_length=4096, hidden_size=256, rotary_percent=0.5
unfused rope: 1.6 ms
fused rope: 0.75 ms

dtype=torch.float32, seq_length=4096, hidden_size=256, rotary_percent=1.0
unfused rope: 2.6 ms
fused rope: 0.58 ms

[crcrpar]

how about using .cuh instead of .h for clarity?

[yaox12]

I just followed the naming convention under csrc/megatron directory, e.g., generic_scaled_masked_softmax.h, scaled_masked_softmax.h, etc.
If you feel .cuh is better, I'm OK to change it.

[crcrpar]

out of curiosity, wouldn't it be useful to have this in apex.transformer.functiona namespace?

[yaox12]

Thanks for your review. I have added two functions to the apex.transformer.functional namespace.

1. fused_apply_rotary_pos_emb, which is a drop-in replacement for the current apply_rotary_pos_emb in Megatron Core.
2. fused_apply_rotary_pos_emb_cached, which would be beneficial when MCore implements caching for the rotary positional embedding.