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

[nvegesna-netizen]

Fixes four silent correctness bugs in Gemma2DotProductAttention that caused sliding window attention (SWA) to either do nothing or produce wrong masks during pretraining and fine-tuning. On top of the corrected unfused path, adds Gemma2FlexDotProductAttention — a fused attention subclass that uses PyTorch FlexAttention (built-in since PyTorch 2.5) to run softcap and SWA in a single Triton kernel, falling back to the fixed unfused parent when a padding mask is present or dropout is active.

Changelog

gemma2_provider.py: Fix SWA gate from if attention_mask is not None and self.window_size is not None: to if self.window_size is not None:. The old gate was never True on MCore's pretrain path because MCore passes attention_mask=None, so SWA was silently skipped for every pretrain step. The mask must always be generated for SWA (even-numbered) layers: these layers use attn_mask_type=arbitrary, and FusedScaleMaskSoftmax with arbitrary + mask=None routes through ScaledSoftmax — plain softmax with no causal masking — so omitting the mask for any sequence length (including short ones where the window covers the full sequence) drops causal masking entirely. get_swa() degenerates to a pure causal mask when the window covers all positions, so always building it is correct.

gemma2_provider.py: Fix AttnMaskType.causal → AttnMaskType.arbitrary for SWA (even-numbered) layers. FusedScaleMaskSoftmax with causal routes through ScaledUpperTriangMaskedSoftmax.apply(input, scale), which takes no mask argument and silently discards the externally generated SWA mask. Switching to arbitrary routes through ScaledMaskedSoftmax, which applies it correctly. Odd-numbered layers remain causal and keep the fast fused path.

gemma2_provider.py: Fix SWA mask dimensionality — get_swa() returns [sq, sk]; the fused CUDA softmax kernel requires a 4D mask. Now unsqueezed to [1, 1, sq, sk] when there is no padding mask, broadcasting correctly to [b, np, sq, sk].

gemma2_provider.py: Fix fine-tuning padding mask combination — SWA mask is now ORed with any incoming padding mask (swa_mask | attention_mask) rather than replacing it. Both masks use True=masked-out; [sq, sk] broadcasts to [b, 1, sq, sk].

gemma2_provider.py: Fix Gemma2ModelProvider.window_size default from (4096, 0) to (4095, 0). window_size[0] counts prior tokens exclusively, so a 4096-token window requires 4095, consistent with the gemma2_bridge.py round-trip convention.

gemma2_provider.py: Replace bare assert statements (stripped by -O) with raise ValueError for CP > 1 and packed-seq guards.

gemma2_provider.py: Add Gemma2FlexDotProductAttention — uses PyTorch FlexAttention (built-in, PyTorch 2.5+) to fuse softcap and SWA into a single Triton kernel via score_mod and block_mask. Module-level _get_softcap_score_mod is lru_cache-decorated so all layers with the same softcap share one function object, avoiding N redundant torch.compile recompilations at startup (torch.compile guards on id(fn)). flex_attention itself is wrapped with torch.compile to trigger the fused Triton path. Falls back to the unfused parent when FlexAttention is unavailable, a padding mask is present, or dropout is active. gemma2_layer_spec updated to use this class as core_attention.

test_gemma2_provider.py: Add TestGemma2DotProductAttention with 7 tests covering all correctness fixes: CP ValueError, packed-seq ValueError, SWA gate fires on pretrain path (attention_mask=None), odd layers
test_gemma2_provider.py: Add TestGemma2DotProductAttention with 7 tests covering all correctness fixes: CP ValueError, packed-seq ValueError, SWA gate fires on pretrain path (attention_mask=None), odd layers have no SWA, even layers use AttnMaskType.arbitrary, SWA combined with padding mask via OR, and window_size default. Add TestGemma2FlexDotProductAttention with 12 tests covering fused path invocation, all fallback conditions (no FlexAttention, dropout, padding mask), output shape, softmax scale, score_mod softcap correctness, lru_cache sharing invariant, block mask caching, SWA layer block mask, and packed-seq ValueError.

Additional Information

Related to:

• MCore fix: Perf configs after refactoring #2357: arbitrary mask type required for externally generated masks
• MCore Cherry pick Malay/dsv3 2509 (669) into r0.1.0 #752: confirms MCore passes attention_mask=None on the pretrain path
• MB PR [bridge] Fix off-by-one in sliding window size for Gemma2, Gemma3, Mistral, and GPT-OSS #2656: companion off-by-one fix in the bridge conversion layer

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[yaoyu-33]

/claude review

[claude]

Bug: division by zero when softcap=0.0

The unfused path guards against zero softcap (if not scale: return logits in logit_softcapping), but this score_mod doesn't. If attn_logit_softcapping is set to 0.0 (to disable softcapping), score / softcap produces inf/nan, and 0.0 * tanh(nan) = nan for any score value of exactly 0.0 (common for padding positions).

Gemma2 always uses 50.0 so this isn't triggered today, but it's a latent correctness bug if anyone disables softcapping on the fused path.

Suggested change

	def _score_mod(score, b, h, q_idx, kv_idx):
	return softcap * torch.tanh(score / softcap)
	def _score_mod(score, b, h, q_idx, kv_idx):
	if softcap == 0.0:
	return score
	return softcap * torch.tanh(score / softcap)

[claude]

Nit: logger.warning for a success message is semantically wrong — this is informational, not a warning condition. Both this line and line 73 should use logger.info.

Suggested change

	logger.warning("Gemma2: PyTorch FlexAttention available — softcap+SWA fused via Triton score_mod.")
	logger.info("Gemma2: PyTorch FlexAttention available — softcap+SWA fused via Triton score_mod.")

[claude]

Light Code Review

Solid set of fixes. The FlexAttention fused path is well-structured with appropriate fallback conditions, and the lru_cache sharing for score_mod is a nice touch to avoid torch.compile recompilation overhead.
Issues:

1. Bug: _get_softcap_score_mod division by zero when softcap=0.0 (inline comment) - The unfused logit_softcapping() function guards against zero scale, but the score_mod closure does not. score / 0.0 produces inf/nan, and 0.0 * tanh(nan) = nan. Gemma2 always uses 50.0 so this is latent, but it will silently produce NaN if anyone disables softcapping on the fused path.
2. Minor: logger.warning for informational messages (inline comment) - Lines 67 and 73 use logger.warning for feature-availability diagnostics. Both should be logger.info.
3. Minor: _build_flex_block_mask ignores the right window component - _flex_window_size[1] (the right window) is stored but never read in _build_flex_block_mask. For Gemma2 where right=0, the causal constraint handles it. But if right > 0 were used, the SWA mask would silently omit the right-window constraint.
   Test Coverage - thorough. Two gaps worth considering: (a) No test for softcap=0.0 on the fused path (would catch the division-by-zero bug above); (b) No test for GQA with FlexAttention (_make_flex_attention uses num_query_groups == num_attention_heads so enable_gqa=True is never exercised).
   Suggested test cases: No perf tests impacted.