> This page is for version 0.5.0 · 26.06.
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.components.attention.utils

## Module Contents

### Functions

| Name                                                                                                                    | Description                                                    |
| ----------------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------- |
| [`initialize_attn_module_and_func`](#nemo_automodel-components-attention-utils-initialize_attn_module_and_func)         | Initialize an attention backend module and callable.           |
| [`postprocess_output_for_attn`](#nemo_automodel-components-attention-utils-postprocess_output_for_attn)                 | Postprocess attention output based on attn\_impl requirements. |
| [`preprocess_args_and_kwargs_for_attn`](#nemo_automodel-components-attention-utils-preprocess_args_and_kwargs_for_attn) | Preprocess attention inputs based on backend requirements.     |

### API

```python
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.

```python
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.

```python
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:**

**`q`** `torch.Tensor`

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

---

**`k`** `torch.Tensor`

Key tensor of shape \[batch, sequence, kv\_heads, head\_dim] or
\[tokens, kv\_heads, head\_dim] for THD input.

---

**`v`** `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`** `Any` — default: \{}

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.