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 :func:`sdpa_fp32_attention_forward` as the process-wide `"sdpa"` implementation.
`sdpa_fp32_attention_forward`	Run SDPA with fp32 query/key/value and cast the output back to the input dtype.

API

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

Install :func: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: typing.Optional[torch.Tensor] = None,
    kwargs = {}
) -> tuple[torch.Tensor, typing.Optional[torch.Tensor]]

Run SDPA with fp32 query/key/value and cast the output back to the input dtype.