DevelopmentAPI ReferenceFull Library ReferenceNemo AutomodelNemo AutomodelComponentsModelsGemma4 Moe

nemo_automodel.components.models.gemma4_moe.sdpa_fp32

View as Markdown

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

Module Contents

Functions

NameDescription
enable_gemma4_sdpa_fp32Install :func:sdpa_fp32_attention_forward as the process-wide "sdpa" implementation.
sdpa_fp32_attention_forwardRun 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.