> 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.models.bagel.attention_masks

Packed multimodal attention mask helpers.

## Module Contents

### Functions

| Name                                                                                                         | Description                                                            |
| ------------------------------------------------------------------------------------------------------------ | ---------------------------------------------------------------------- |
| [`_repeat_document_labels`](#nemo_automodel-components-models-bagel-attention_masks-_repeat_document_labels) | Return a per-token document id for packed samples.                     |
| [`_repeat_split_labels`](#nemo_automodel-components-models-bagel-attention_masks-_repeat_split_labels)       | Return per-token split labels used by the packed block-mask predicate. |
| [`create_sparse_mask`](#nemo_automodel-components-models-bagel-attention_masks-create_sparse_mask)           | Create the block-mask predicate for packed multimodal attention.       |

### API

```python
nemo_automodel.components.models.bagel.attention_masks._repeat_document_labels(
    document_lens: list[int],
    device: torch.device
) -> torch.Tensor
```

Return a per-token document id for packed samples.

```python
nemo_automodel.components.models.bagel.attention_masks._repeat_split_labels(
    split_lens: list[int],
    attn_modes: list[str],
    device: torch.device
) -> tuple[torch.Tensor, torch.Tensor]
```

Return per-token split labels used by the packed block-mask predicate.

```python
nemo_automodel.components.models.bagel.attention_masks.create_sparse_mask(
    document_lens: list[int],
    split_lens: list[int],
    attn_modes: list[str],
    device: torch.device
) -> typing.Any
```

Create the block-mask predicate for packed multimodal attention.