> This page is for version 0.4.0 · 26.04.
> 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.models.step5.flops

Effective Step-5 text FLOPs, including CSA's forward-only discrete indexer.

## Module Contents

### Functions

| Name                                                                             | Description                                                                       |
| -------------------------------------------------------------------------------- | --------------------------------------------------------------------------------- |
| [`_document_work`](#nemo_automodel-components-models-step5-flops-_document_work) | Count independent documents; each length excludes physical pack padding.          |
| [`_sum_history`](#nemo_automodel-components-models-step5-flops-_sum_history)     | Sum complete historical regions for document-local queries zero through length-1. |
| [`step5_flops`](#nemo_automodel-components-models-step5-flops-step5_flops)       | Count Step-5 text training GEMMs for independent, fully supervised documents.     |

### API

```python
nemo_automodel.components.models.step5.flops._document_work(
    config: nemo_automodel.components.models.step5.configuration.Step5Config,
    lengths: list[int],
    supervised_tokens: int
) -> float
```

Count independent documents; each length excludes physical pack padding.

```python
nemo_automodel.components.models.step5.flops._sum_history(
    length: int,
    region_size: int,
    cap: int | None
) -> int
```

Sum complete historical regions for document-local queries zero through length-1.

```python
nemo_automodel.components.models.step5.flops.step5_flops(
    config: nemo_automodel.components.models.step5.configuration.Step5Config,
    gbs: int = 1,
    seq_len: int | None = None
) -> float
```

Count Step-5 text training GEMMs for independent, fully supervised documents.

The count includes trainable linear/attention forward and backward (6 x MACs)
and the forward-only indexer (2 x MACs). Embedding lookup, norms, activations,
rotary/softmax elementwise work, vision, MTP, optimizer work, kernel padding,
masked indexer scores and activation recomputation are excluded. A CSA query
at position t sees region_size \* min(topk, floor(t / region_size)) historical
tokens plus t % region_size + 1 tokens in its current region.

**Parameters:**

**`config`** `Step5Config`

Actual complete text configuration, before pipeline layer partitioning.

---

**`gbs`** `int` — default: 1

Number of independent documents.

---

**`seq_len`** `int | None` — default: None

Tokens per document; defaults to max_position_embeddings.

---

**Returns:** `float`

Effective training FLOPs with all document tokens supervised.