[main] fix(moe): fix the bug where gate was not sliced when kv_head < tp_size. - #3575

[main] fix(moe): fix the bug where gate was not sliced when kv_head < tp_size.#3575

Phlip79 merged 5 commits into

NVIDIA:mainNVIDIA/Megatron-LM:mainfrom

yuzhongw-nvidia:yuzhongw/fix_output_gateyuzhongw-nvidia/Megatron-LM:yuzhongw/fix_output_gateCopy head branch name to clipboard