nemo_automodel.components.models.gemma4_moe.attention

Composition of document-causal masks with Gemma4’s native attention policy.

Module Contents

Functions

Name	Description
`prepare_gemma4_document_mask`	Compose native full/sliding attention with an explicit document-causal mask.

API

nemo_automodel.components.models.gemma4_moe.attention.prepare_gemma4_document_mask(
    config: transformers.configuration_utils.PreTrainedConfig,
    attention_mask: torch.Tensor,
    mm_token_type_ids: torch.Tensor | None,
    document_ids: torch.Tensor,
    dtype: torch.dtype
) -> dict[str, torch.Tensor]

Compose native full/sliding attention with an explicit document-causal mask.

Call only for the document-causal contract marked by _packed_seq_ids; arbitrary custom 4D masks are already complete and must not be interpreted here. HF otherwise treats a 4D input as complete, bypassing its vision/window policy.

Parameters:

config

PreTrainedConfig

Gemma4 text configuration using eager or SDPA attention.

attention_mask

torch.Tensor

Boolean-allowed or additive-zero/minimum mask [batch, 1, sequence, sequence]. This carries the caller’s selected document isolation and excludes padding/fake media on both axes.

mm_token_type_ids

torch.Tensor | None

Processor modality IDs [batch, sequence], or None for text.

document_ids

torch.Tensor

Original document IDs [batch, sequence]. These keep image groups separate even when the selected mask allows cross-document attention.

dtype

torch.dtype

Model floating point dtype for the additive native masks.

Returns: dict[str, torch.Tensor]

Full and sliding masks [batch, 1, sequence, sequence]: boolean for SDPA