> This page is for version Latest.
> 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.7.0 · 26.10: https://docs.nvidia.com/nemo/automodel/v0.7/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.

# nemo_automodel.components.datasets.zephon.attention

Stateless mapping from Zephon tensors to AutoModel's packed input contract.

## Module Contents

### Classes

| Name                                                                                  | Description                                                             |
| ------------------------------------------------------------------------------------- | ----------------------------------------------------------------------- |
| [`TrainingBatch`](#nemo_automodel-components-datasets-zephon-attention-TrainingBatch) | Collated CPU tensors; AutoModel aliases are applied only when yielding. |

### Functions

| Name                                                                                                                            | Description                                                           |
| ------------------------------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------- |
| [`adapt_model_inputs`](#nemo_automodel-components-datasets-zephon-attention-adapt_model_inputs)                                 | Prepare attention membership from Zephon's collated document IDs.     |
| [`enable_xdoc_attention`](#nemo_automodel-components-datasets-zephon-attention-enable_xdoc_attention)                           | Join attention and position ranges inside each packed batch document. |
| [`require_supported_packing_contract`](#nemo_automodel-components-datasets-zephon-attention-require_supported_packing_contract) | Validate model input interfaces once, during loader construction.     |

### API

```python
class nemo_automodel.components.datasets.zephon.attention.TrainingBatch
```

**Bases:** `typing.TypedDict`

Collated CPU tensors; AutoModel aliases are applied only when yielding.

Token fields have shape \[batch, sequence]. `document_ids` preserves
subdocument membership (zero for padding); `attention_document_ids` also
excludes processor-ineligible tokens and joins subdocuments when xdoc is on.
Model positions are \[batch, sequence] or \[axes, batch, sequence] for mRoPE.

**`_packed_seq_ids`** `NotRequired[Tensor]`

---

**`attention_document_ids`** `NotRequired[Tensor]`

---

**`attention_mask`** `NotRequired[Tensor]`

---

**`cu_seqlens`** `NotRequired[Tensor]`

---

**`document_ids`** `NotRequired[Tensor]`

---

**`input_ids`** `Tensor`

---

**`labels`** `Tensor`

---

**`max_seqlen`** `NotRequired[int]`

---

**`packed_token_indices`** `NotRequired[Tensor]`

---

**`padding_mask`** `Tensor`

---

**`position_ids`** `Tensor`

---

**`qkv_format`** `NotRequired[str]`

---

**`seq_lens`** `NotRequired[Tensor]`

---

**`seq_lens_padded`** `NotRequired[Tensor]`

---

```python
nemo_automodel.components.datasets.zephon.attention.adapt_model_inputs(
    batch: nemo_automodel.components.datasets.zephon.attention.TrainingBatch,
    attention_eligibility: torch.Tensor | None = None
) -> nemo_automodel.components.datasets.zephon.attention.TrainingBatch
```

Prepare attention membership from Zephon's collated document IDs.

This follows to_training so IDs already match input tokens after shifting and
padding. Zephon owns structural membership; processor eligibility and the
consumer's attention policy belong to this integration. AutoModel field names
are applied later, when the loader yields the completed batch.

**Parameters:**

**`batch`** `TrainingBatch`

Owned mapping with int64 `document_ids` \[batch, sequence] from
Zephon. Token/label/padding tensors have the same shape; model positions
may additionally have a leading RoPE-axes dimension. Mutates mapping
entries only; never modifies an existing tensor in place.

---

**`attention_eligibility`** `torch.Tensor | None` — default: None

Optional bool \[batch, sequence], supplied by the
processor to exclude synthetic image tokens independently of padding.

---

**Returns:** `TrainingBatch`

The same mapping with `attention_document_ids` \[batch, sequence]. It

```python
nemo_automodel.components.datasets.zephon.attention.enable_xdoc_attention(
    batch: nemo_automodel.components.datasets.zephon.attention.TrainingBatch,
    packing_contract: nemo_automodel.components.datasets.packing.PackedSequenceContract
) -> nemo_automodel.components.datasets.zephon.attention.TrainingBatch
```

Join attention and position ranges inside each packed batch document.

**Parameters:**

**`batch`** `TrainingBatch`

Mapped training tensors as documented by `adapt_model_inputs`;
`document_ids` and `attention_document_ids` have shape
\[batch, sequence]; zero in the latter marks excluded positions. Model
`position_ids` has shape \[batch, sequence] or
\[position_axes, batch, sequence] for multi-axis RoPE. Subdocument
coordinates are nonnegative and local to each subdocument. The owned
mapping is updated in place, but no existing tensor is mutated.

---

**`packing_contract`** `PackedSequenceContract`

Model-owned packed attention representation.

---

**Returns:** `TrainingBatch`

The same mapping with translated position IDs and joined attention IDs

```python
nemo_automodel.components.datasets.zephon.attention.require_supported_packing_contract(
    contract: nemo_automodel.components.datasets.packing.PackedSequenceContract
) -> None
```

Validate model input interfaces once, during loader construction.

The text and image config builders call this before sources/workers start;
the contract is fixed for that loader and is not rechecked for each batch.

`return_packed_sequence_metadata` specifically requests AutoModel's
gathered-token bundle: packed_token_indices \[batch, sequence] (-1 at padding),
cu_seqlens \[batch, boundaries] (-1 at unused entries), and max_seqlen. These
describe the stream after gathering non-padding tokens. AutoModel's FA4 and
recurrent/hybrid models use it for kernel layouts and recurrent-state resets.
Zephon's cu_seqlens instead describes physical offsets including padding.
The final batch map computes the gathered-token offsets with AutoModel's
`build_packed_sequence_metadata` on the selected attention document IDs.
Xdoc therefore joins both attention and recurrent-state boundaries; MTP
uses those same selected boundaries. Keeping this bundle batch-major
lets the model flatten it after PP microbatch splitting.