nemo_automodel.components.attention.utils

Module Contents

Functions

Name	Description
`initialize_attn_module_and_func`	Initialize an attention backend module and callable.
`postprocess_output_for_attn`	Postprocess attention output based on attn_impl requirements.
`preprocess_args_and_kwargs_for_attn`	Preprocess attention inputs based on backend requirements.

API

nemo_automodel.components.attention.utils.initialize_attn_module_and_func(
    attn_impl: str,
    num_attention_heads: int,
    num_qk_channels: int,
    num_v_channels: int,
    softmax_scale: float,
    attn_mask_type: str = 'causal',
    qkv_format: str = 'bshd',
    num_gqa_groups: int | None = None,
    kwargs: typing.Any = {}
) -> tuple[torch.nn.Module | None, typing.Callable[..., torch.Tensor]]

Initialize an attention backend module and callable.

nemo_automodel.components.attention.utils.postprocess_output_for_attn(
    x: torch.Tensor,
    attn_impl: str
) -> torch.Tensor

Postprocess attention output based on attn_impl requirements.

nemo_automodel.components.attention.utils.preprocess_args_and_kwargs_for_attn(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    attention_mask: torch.Tensor | None,
    attn_impl: str,
    kwargs: typing.Any = {}
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, dict[str, typing.Any]]

Preprocess attention inputs based on backend requirements.

Parameters:

torch.Tensor

Query tensor of shape [batch, sequence, heads, head_dim] or [tokens, heads, head_dim] for THD input.

torch.Tensor

Key tensor of shape [batch, sequence, kv_heads, head_dim] or [tokens, kv_heads, head_dim] for THD input.

torch.Tensor

Value tensor with the same leading layout as k and an optional distinct value head dimension.

attention_mask

torch.Tensor | None

Optional tensor of shape [batch, sequence] for padding or indexed packing, or [batch, 1, sequence, sequence] for an explicit dense mask. TE accepts only a binary [batch, sequence] padding mask (1/True for valid tokens, 0/False for padding). Packed TE inputs require None and THD sequence metadata instead.

attn_impl

str

Attention backend name.

**kwargs

AnyDefaults to {}

Backend metadata. Packed FA4 accepts cu_seqlens of shape [documents + 1] and packed_token_indices of shape [tokens]. For SDPA and FA4, a positive left window_size counts visible keys including the current token; negative/None means unbounded. The legacy TE branch forwards its backend-specific distances unchanged. Packed TE accepts cu_seqlens of shape [documents + 1] and optional cu_seqlens_padded with the same shape, both int32, together with an integer max_seqlen.

Returns: torch.Tensor

Query, key, and value tensors in the backend layout plus its keyword

Raises:

ValueError: If the selected backend cannot consume the input layout.