> 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.datasets.packed_seq

Typed packed-sequence metadata for neat sequence packing.

NeMo AutoModel's neat packing encodes document boundaries in an indexed map
`[B, S]` where each position holds the 1-based document index it belongs to
(`0` = padding), e.g. `[1, 1, 2, 2, 2, 0]`. This module turns that map into
typed metadata that drives HuggingFace's varlen FlashAttention path via
`transformers.modeling_flash_attention_utils.FlashAttentionKwargs`
(`cu_seq_lens_q`/`cu_seq_lens_k`/`max_length_q`/`max_length_k`).

The varlen path reshapes `[B, S, ...]` to `[B * S, ...]` without unpadding
(see the `is_fa_with_varlen_kwargs` branch in
`transformers.modeling_flash_attention_utils._flash_attention_forward`).
`cu_seqlens` therefore spans the whole flattened batch, including padding runs,
and a boundary is forced at every row start so a document that fills one row
cannot merge with the first document of the next row once the batch is flattened.

## Module Contents

### Classes

| Name                                                                                | Description                                                        |
| ----------------------------------------------------------------------------------- | ------------------------------------------------------------------ |
| [`PackedSeqParams`](#nemo_automodel-components-datasets-packed_seq-PackedSeqParams) | Packed-sequence metadata derived from a neat-packing document map. |

### Functions

| Name                                                                                                              | Description                                                           |
| ----------------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------- |
| [`packed_seq_params_from_doc_ids`](#nemo_automodel-components-datasets-packed_seq-packed_seq_params_from_doc_ids) | Build packed-sequence metadata from a neat-packing document map.      |
| [`to_flash_attention_kwargs`](#nemo_automodel-components-datasets-packed_seq-to_flash_attention_kwargs)           | Convert packed-sequence params to HuggingFace `FlashAttentionKwargs`. |

### API

```python
class nemo_automodel.components.datasets.packed_seq.PackedSeqParams(
    cu_seqlens: torch.Tensor,
    max_seqlen: int,
    doc_ids: torch.Tensor
)
```

Dataclass

Packed-sequence metadata derived from a neat-packing document map.

**`cu_seqlens`** `Tensor`

Int32 tensor of shape `[num_segments + 1]`. Cumulative token
counts over the flattened `[batch * sequence]` axis, including
trailing padding runs, so `cu_seqlens[-1] == batch * sequence`.
Segment `k` spans flattened positions
`[cu_seqlens[k], cu_seqlens[k + 1])`.

---

**`doc_ids`** `Tensor`

Int tensor of shape `[batch, sequence]` with the 1-based
document index per token (`0` = padding). Preserved unchanged for
loss functions and context-parallel consumers.

---

**`max_seqlen`** `int`

Length of the longest segment, sized for the varlen
flash-attention kernel.

---

```python
nemo_automodel.components.datasets.packed_seq.packed_seq_params_from_doc_ids(
    doc_ids: torch.Tensor
) -> nemo_automodel.components.datasets.packed_seq.PackedSeqParams
```

Build packed-sequence metadata from a neat-packing document map.

Segments are maximal runs of an identical document index within a row, plus a
forced boundary at every row start. Padding runs (index `0`) become their
own segments so `cu_seqlens` covers the full flattened `[batch * sequence]`
axis that the varlen flash-attention path consumes without unpadding. Padding
tokens therefore attend only among themselves, which is harmless because their
labels are ignored by the loss.

**Parameters:**

**`doc_ids`** `torch.Tensor`

Int tensor of shape `[batch, sequence]` where each position
holds the 1-based document index it belongs to (`0` = padding).

---

**Returns:** `PackedSeqParams`

PackedSeqParams carrying `cu_seqlens`, `max_seqlen`, and the original

```python
nemo_automodel.components.datasets.packed_seq.to_flash_attention_kwargs(
    params: nemo_automodel.components.datasets.packed_seq.PackedSeqParams
) -> dict[str, torch.Tensor | int]
```

Convert packed-sequence params to HuggingFace `FlashAttentionKwargs`.

The returned mapping matches `transformers` `FlashAttentionKwargs` so a
packed batch drives `flash_attn_varlen_func` directly. Query and key share
one layout (self-attention), so the `q` and `k` entries are equal.

**Parameters:**

**`params`** `PackedSeqParams`

Packed-sequence metadata. `params.cu_seqlens` is an int32 tensor
of shape `[num_segments + 1]`.

---

**Returns:** `dict[str, torch.Tensor | int]`

Mapping with `cu_seq_lens_q`/`cu_seq_lens_k` (the int32 `cu_seqlens`