fix+feat(gemma2): fix SWA correctness bugs and add FlexAttention fused softcap+SWA path#4308

nvegesna-netizen wants to merge 5 commits into

NVIDIA-NeMo:mainNVIDIA-NeMo/Megatron-Bridge:mainfrom

nvegesna-netizen:nvegesna/fix-gemma2-swa-and-assertionsnvegesna-netizen/Megatron-Bridge:nvegesna/fix-gemma2-swa-and-assertionsCopy head branch name to clipboard