nemo_automodel.components.models.gemma4_moe.sdpa_fp32

fp32 SDPA attention for Gemma4 (fix for #2208: fused bf16 SDPA NaNs on Hopper).

Module Contents

Functions

Name	Description
`enable_gemma4_sdpa_fp32`	Install `sdpa_fp32_attention_forward` as the process-wide `"sdpa"` implementation.
`sdpa_fp32_attention_forward`	Run SDPA in FP32 and cast the output back to the query dtype.

API

nemo_automodel.components.models.gemma4_moe.sdpa_fp32.enable_gemma4_sdpa_fp32() -> None

Install sdpa_fp32_attention_forward as the process-wide "sdpa" implementation.

nemo_automodel.components.models.gemma4_moe.sdpa_fp32.sdpa_fp32_attention_forward(
    module: torch.nn.Module,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
    kwargs = {}
) -> tuple[torch.Tensor, torch.Tensor | None]

Run SDPA in FP32 and cast the output back to the query dtype.

Parameters:

module

torch.nn.Module

Attention module holding scaling and causal-mask settings.

query

torch.Tensor

Tensor of shape [batch, query_heads, queries, head_dim].

key

torch.Tensor

Tensor of shape [batch, kv_heads, keys, head_dim].

value

torch.Tensor

Tensor with key’s layout.

attention_mask

torch.Tensor | NoneDefaults to None

Optional mask broadcastable to [batch, query_heads, queries, keys].

**kwargs

Defaults to {}

Arguments forwarded to the Hugging Face SDPA implementation.

Returns: torch.Tensor

Attention output [batch, queries, query_heads, head_dim] in query’s