[BUG]megatron-lm，with torchompile，The provided qkv memory layout is not supported!

[qingshanxwx]

code: the main branch https://github.com/NVIDIA/Megatron-LM
transformer_engine 1.13
eager is normal, compile will report The provided qkv memory layout is not supported!
The error is as follows：

According to the functional description, the stride of k and v should be (SHD, HD, D, 1)， but it is actually (HD, H*D, D, 1)

megatron/core/extensions/transformer_engine.py

printed result

Then call DotProductAttention, v'stride will become (524288, 128, 32, 1) directly，k'stride not change.

printed result