Fix CP crash with GQA + asymmetric KV head dims (#2868)#2901

beccohov wants to merge 1 commit into

NVIDIA:mainNVIDIA/TransformerEngine:mainfrom

beccohov:bugfix/cp-gqa-asymmetric-kv-head-dimsbeccohov/TransformerEngine:bugfix/cp-gqa-asymmetric-kv-head-dimsCopy head branch name to clipboard