> 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.utils.flops_utils

## Module Contents

### Functions

| Name                                                                                                              | Description                                                                                             |
| ----------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------- |
| [`_build_moe_layer_pattern`](#nemo_automodel-components-utils-flops_utils-_build_moe_layer_pattern)               | Build a list of 0/1 indicating dense(0) vs MoE(1) per layer.                                            |
| [`_gdn_attention_per_layer_flops`](#nemo_automodel-components-utils-flops_utils-_gdn_attention_per_layer_flops)   | FLOPs for a single Gated DeltaNet (GDN / linear attention) layer.                                       |
| [`_hybrid_model_flops`](#nemo_automodel-components-utils-flops_utils-_hybrid_model_flops)                         | Model FLOPs for hybrid model                                                                            |
| [`_kda_attention_per_layer_flops`](#nemo_automodel-components-utils-flops_utils-_kda_attention_per_layer_flops)   | Per-layer FLOPs for Kimi Delta Attention (KDA, gated delta-rule linear attention).                      |
| [`_mamba_layer_flops`](#nemo_automodel-components-utils-flops_utils-_mamba_layer_flops)                           | Model FLOPs for Mamba layer.                                                                            |
| [`_mla_attention_per_layer_flops`](#nemo_automodel-components-utils-flops_utils-_mla_attention_per_layer_flops)   | Per-layer FLOPs for Multi-Latent Attention (MLA).                                                       |
| [`_mla_moe_model_flops`](#nemo_automodel-components-utils-flops_utils-_mla_moe_model_flops)                       | FLOPs for MLA + MoE transformer models (DeepSeek-V3 style).                                             |
| [`_nemotronh_mlp_layer_flops`](#nemo_automodel-components-utils-flops_utils-_nemotronh_mlp_layer_flops)           | Model FLOPs for MLP layer. Assume gated linear unit.                                                    |
| [`_nemotronh_moe_layer_flops`](#nemo_automodel-components-utils-flops_utils-_nemotronh_moe_layer_flops)           | Model FLOPs for a MoE layer in Nemotron V3/Super V3 (hybrid Mamba/Attention/MoE).                       |
| [`_nemotronh_mtp_flops`](#nemo_automodel-components-utils-flops_utils-_nemotronh_mtp_flops)                       | Model FLOPs for the Multi-Token-Prediction (MTP) head of Nemotron-3 Super/Ultra.                        |
| [`_non_mla_attn_layer_flops`](#nemo_automodel-components-utils-flops_utils-_non_mla_attn_layer_flops)             | Model FLOPs for attention layer                                                                         |
| [`_sum_min_floor_div`](#nemo_automodel-components-utils-flops_utils-_sum_min_floor_div)                           | `sum_{j=1..seq_len} min(cap, floor(j / ratio))` in closed form (`cap=None` disables the cap). |
| [`_sum_min_window`](#nemo_automodel-components-utils-flops_utils-_sum_min_window)                                 | `sum_{i=0..seq_len-1} min(i + 1, window)`: causal sliding-window keys per query, summed.      |
| [`_sum_mod`](#nemo_automodel-components-utils-flops_utils-_sum_mod)                                               | `sum_{j=1..seq_len} (j mod ratio)` in closed form.                                            |
| [`attention_flops_calculator`](#nemo_automodel-components-utils-flops_utils-attention_flops_calculator)           | Calculate the flops for the attention part.                                                             |
| [`bert_flops`](#nemo_automodel-components-utils-flops_utils-bert_flops)                                           | Model FLOPs for BERT family - accepts either AutoConfig or normalized config                            |
| [`calculate_mfu`](#nemo_automodel-components-utils-flops_utils-calculate_mfu)                                     | Calculate Model FLOPs Utilization (MFU).                                                                |
| [`clip_vit_l_flops`](#nemo_automodel-components-utils-flops_utils-clip_vit_l_flops)                               | Model FLOPs for CLIP ViT                                                                                |
| [`deepseek_v41_flops`](#nemo_automodel-components-utils-flops_utils-deepseek_v41_flops)                           | Model FLOPs for DeepSeek-V4.1 (CSA2 sparse attention, single-pass mHC, Engram, MoE).                    |
| [`deepseekv3_flops`](#nemo_automodel-components-utils-flops_utils-deepseekv3_flops)                               | Model FLOPs for DeepSeek V3 - accepts either AutoConfig or normalized config                            |
| [`flux_flops`](#nemo_automodel-components-utils-flops_utils-flux_flops)                                           | Model FLOPs for FLUX                                                                                    |
| [`get_flops_formula_for_hf_config`](#nemo_automodel-components-utils-flops_utils-get_flops_formula_for_hf_config) | Get the appropriate FLOPs formula function for a given HuggingFace config.                              |
| [`glm4_moe_flops`](#nemo_automodel-components-utils-flops_utils-glm4_moe_flops)                                   | Estimate FLOPs for GLM4 MoE model configurations.                                                       |
| [`gpt3_flops`](#nemo_automodel-components-utils-flops_utils-gpt3_flops)                                           | Model FLOPs for GPT3 family - accepts either AutoConfig or normalized config                            |
| [`gpt_oss_flops`](#nemo_automodel-components-utils-flops_utils-gpt_oss_flops)                                     | Model FLOPs for GPT-OSS                                                                                 |
| [`gpt_oss_flops_calculator`](#nemo_automodel-components-utils-flops_utils-gpt_oss_flops_calculator)               | Calculate the flops for the GPT-OSS model                                                               |
| [`kimi_k3_flops`](#nemo_automodel-components-utils-flops_utils-kimi_k3_flops)                                     | Model FLOPs for Kimi K3 (hybrid KDA / gated-MLA attention + latent MoE with a SiTU shared expert).      |
| [`llama2_flops`](#nemo_automodel-components-utils-flops_utils-llama2_flops)                                       | Model FLOPs for llama2 family - accepts either AutoConfig or normalized config                          |
| [`llama3_flops`](#nemo_automodel-components-utils-flops_utils-llama3_flops)                                       | Model FLOPs for llama3 family - accepts either AutoConfig or normalized config                          |
| [`loss_flops_calculator`](#nemo_automodel-components-utils-flops_utils-loss_flops_calculator)                     | Calculate the flops for the loss                                                                        |
| [`minimax_m2_flops`](#nemo_automodel-components-utils-flops_utils-minimax_m2_flops)                               | Model FLOPs for MiniMax-M2 family - accepts either AutoConfig or normalized config.                     |
| [`mixtral_flops`](#nemo_automodel-components-utils-flops_utils-mixtral_flops)                                     | Model FLOPs for mixtral family - accepts either AutoConfig or normalized config                         |
| [`mla_moe_flops`](#nemo_automodel-components-utils-flops_utils-mla_moe_flops)                                     | Model FLOPs for MLA + MoE models (Kimi K2, GLM-5, Mistral Small 4, etc.).                               |
| [`moe_mlp_flops_calculator`](#nemo_automodel-components-utils-flops_utils-moe_mlp_flops_calculator)               | Calculate the flops for the MLP                                                                         |
| [`nemotron_flops`](#nemo_automodel-components-utils-flops_utils-nemotron_flops)                                   | Model FLOPs for nemotron family - accepts either AutoConfig or normalized config                        |
| [`nemotronh_flops`](#nemo_automodel-components-utils-flops_utils-nemotronh_flops)                                 | Model FLOPs for NemotronH                                                                               |
| [`neva_projection_flops`](#nemo_automodel-components-utils-flops_utils-neva_projection_flops)                     | Model FLOPs for NeVA Projection                                                                         |
| [`qwen3_5_flops`](#nemo_automodel-components-utils-flops_utils-qwen3_5_flops)                                     | Model FLOPs for Qwen3.5 family (MoE and Dense) with hybrid GDN/full attention.                          |
| [`qwen3_8_flash_next_flops`](#nemo_automodel-components-utils-flops_utils-qwen3_8_flash_next_flops)               | Model FLOPs for Qwen3.8-Flash-Next (GDN + compressed-block QSA, HyperConnections, Engram PLE, MoE).     |
| [`qwen3_flops`](#nemo_automodel-components-utils-flops_utils-qwen3_flops)                                         | Model FLOPs for Qwen3 family - accepts either AutoConfig or normalized config                           |
| [`step3_5_flash_flops`](#nemo_automodel-components-utils-flops_utils-step3_5_flash_flops)                         | Model FLOPs for Step3.5-Flash (GQA + sliding-window / full attention + MoE).                            |
| [`transformer_flops`](#nemo_automodel-components-utils-flops_utils-transformer_flops)                             | Calculate FLOPs for a standard Transformer model - accepts either AutoConfig or normalized config.      |

### API

```python
nemo_automodel.components.utils.flops_utils._build_moe_layer_pattern(
    config,
    layers
)
```

Build a list of 0/1 indicating dense(0) vs MoE(1) per layer.

Handles multiple config styles: first\_k\_dense\_replace + moe\_layer\_freq,
mlp\_layer\_types list, etc.

```python
nemo_automodel.components.utils.flops_utils._gdn_attention_per_layer_flops(
    gbs,
    seq_len,
    hidden_size,
    linear_key_head_dim,
    linear_value_head_dim,
    linear_num_key_heads,
    linear_num_value_heads,
    linear_conv_kernel_dim
)
```

FLOPs for a single Gated DeltaNet (GDN / linear attention) layer.

Based on the GDN FLOPs calculator from Megatron-Bridge PR #2925.

```python
nemo_automodel.components.utils.flops_utils._hybrid_model_flops(
    config,
    gbs,
    seq_len
)
```

Model FLOPs for hybrid model

```python
nemo_automodel.components.utils.flops_utils._kda_attention_per_layer_flops(
    gbs: int,
    seq_len: int,
    hs: int,
    head_dim: int,
    num_heads: int,
    conv_kernel_size: int,
    use_full_rank_gate: bool,
    chunk_size: int = 64
) -> float
```

Per-layer FLOPs for Kimi Delta Attention (KDA, gated delta-rule linear attention).

Projections follow the KDA layer definition (Kimi Linear, arXiv:2510.26692):
q/k/v with short causal convolutions, the low-rank forget gate
(`f_a`/`f_b`), the beta projection, the output gate (full-rank `g_proj`
or low-rank `g_a`/`g_b`) and `o_proj`. The chunkwise kernel is costed
with the paper's own count, `FLOPs_KDA(T; C, d_h) = 6 T d_h^2 + 3 T C d_h + T C^2`
per head (forward, chunk size `C`), halved to MAC-equivalents so the shared
`6 x` (forward + backward) factor of this module applies.

```python
nemo_automodel.components.utils.flops_utils._mamba_layer_flops(
    config,
    gbs,
    seq_len
)
```

Model FLOPs for Mamba layer.

Three components:

* in\_proj:  input projections (x\_proj, z\_proj, dt\_proj, B\_proj, C\_proj)
* scan:     SSM scan kernel (7x factor accounts for the full SSD scan cost)
* out\_proj: output projection back to hidden\_size
  Multiplied by 6 (3x fwd+bwd \* 2x FMA) for in\_proj/out\_proj (standard GEMMs),
  and 7 \* 3 = 21 for scan (non-GEMM kernel, higher op count per element).

```python
nemo_automodel.components.utils.flops_utils._mla_attention_per_layer_flops(
    gbs,
    seq_len,
    hs,
    attention_heads,
    q_lora_rank,
    kv_lora_rank,
    qk_rope_head_dim,
    qk_nope_head_dim,
    v_head_dim,
    index_topk = None,
    index_n_heads = 0,
    index_head_dim = 0
)
```

Per-layer FLOPs for Multi-Latent Attention (MLA).

Shared by DeepSeek V3, Kimi K2.5, Mistral Small 4, GLM-5, etc.

When index\_topk is set (DSA / sparse attention), accounts for:

* Sparse main attention BMM: S \* index\_topk instead of 0.5 \* S^2
* DSA indexer overhead: Q/K/weights projections + full S^2 indexer BMM

```python
nemo_automodel.components.utils.flops_utils._mla_moe_model_flops(
    gbs,
    seq_len,
    hs,
    layers,
    attention_heads,
    vocab_size,
    q_lora_rank,
    kv_lora_rank,
    qk_rope_head_dim,
    qk_nope_head_dim,
    v_head_dim,
    dense_ffn_hs,
    moe_ffn_hs,
    moe_router_topk,
    moe_shared_expert_hs,
    moe_layer_pattern,
    mtp_num_layers = 0,
    index_topk = None,
    index_n_heads = 0,
    index_head_dim = 0
)
```

FLOPs for MLA + MoE transformer models (DeepSeek-V3 style).

**Parameters:**

**`moe_layer_pattern`**

List of 0/1 per layer (0=dense, 1=MoE).

---

**`moe_shared_expert_hs`**

Total intermediate size for all shared experts combined.

---

**`index_topk`** — default: None

If set, use DSA sparse attention with this many selected positions.

---

**`index_n_heads`** — default: 0

Number of heads in the DSA indexer.

---

**`index_head_dim`** — default: 0

Head dimension of the DSA indexer.

---

```python
nemo_automodel.components.utils.flops_utils._nemotronh_mlp_layer_flops(
    config,
    gbs,
    seq_len
)
```

Model FLOPs for MLP layer. Assume gated linear unit.

```python
nemo_automodel.components.utils.flops_utils._nemotronh_moe_layer_flops(
    config,
    gbs,
    seq_len
)
```

Model FLOPs for a MoE layer in Nemotron V3/Super V3 (hybrid Mamba/Attention/MoE).

Nemotron V3 uses relu2 (non-gated) for both routed and shared experts,
so each expert has 2 linear projections (up\_proj + down\_proj), not 3.

When moe\_latent\_size is set (Super V3), routed experts operate in a reduced
latent space with additional projection layers (fc1\_latent\_proj, fc2\_latent\_proj).
The shared expert and gate always operate in the full hidden\_size dimension.

Accounts for:

1. Routed experts: only num\_experts\_per\_tok activated per token.
2. Shared expert: always active for every token (full hidden\_size).
3. Router/gate: linear projection hidden\_size -> n\_routed\_experts.
4. Latent projections (if moe\_latent\_size is set): down and up projections.

```python
nemo_automodel.components.utils.flops_utils._nemotronh_mtp_flops(
    config,
    gbs,
    seq_len,
    num_mtp_layers,
    mtp_block_types,
    use_repeated_layer
)
```

Model FLOPs for the Multi-Token-Prediction (MTP) head of Nemotron-3 Super/Ultra.

The head predicts `num_mtp_layers` (N) additional tokens. Each of the N depths runs
the MTP block pattern once, plus a depth-fusion projection (`eh_proj`: cat\[enorm(embed),
hnorm(hidden)] of size 2\*hidden -> hidden) and a vocab projection (weight-tied lm\_head).

Repeated layer: when `use_repeated_layer` is True the model builds a SINGLE physical
depth (`mtp_block_types` lists its sublayers) and reuses it across all N depths -- but
it still EXECUTES once per depth, so the block compute is N x the physical sublayers.
That N x is the repeated layer's FLOPs. When False, `mtp_block_types` already spans all
N physical depths, so the block runs once.

These settings are NOT fully recoverable from the HF config (it retains only the physical
depth count and omits the block pattern), so callers pass the effective values read from
the built model: `num_mtp_layers = model.mtp_config.num_layers` and
`mtp_block_types = [s.block_type for s in model.mtp.layers]`.

**Parameters:**

**`num_mtp_layers`**

effective MTP depths actually run (model.mtp\_config.num\_layers).

---

**`mtp_block_types`**

block types ("mamba"/"attention"/"mlp"/"moe") of the physical MTP
sublayers (model.mtp.layers).

---

**`use_repeated_layer`**

True if the physical depth is reused across the N depths.

---

```python
nemo_automodel.components.utils.flops_utils._non_mla_attn_layer_flops(
    config,
    gbs,
    seq_len
)
```

Model FLOPs for attention layer

```python
nemo_automodel.components.utils.flops_utils._sum_min_floor_div(
    seq_len: int,
    ratio: int,
    cap: int | None
) -> int
```

`sum_{j=1..seq_len} min(cap, floor(j / ratio))` in closed form (`cap=None` disables the cap).

`floor(j / ratio)` is the number of complete compressed groups visible to the query at
zero-based position `j - 1`; the cap is the sparse top-k. Closed form so million-token
sequences cost nothing to evaluate.

```python
nemo_automodel.components.utils.flops_utils._sum_min_window(
    seq_len: int,
    window: int
) -> int
```

`sum_{i=0..seq_len-1} min(i + 1, window)`: causal sliding-window keys per query, summed.

```python
nemo_automodel.components.utils.flops_utils._sum_mod(
    seq_len: int,
    ratio: int
) -> int
```

`sum_{j=1..seq_len} (j mod ratio)` in closed form.

`j mod ratio` is the length of the incomplete compressed tail visible to the query at
zero-based position `j - 1`.

```python
nemo_automodel.components.utils.flops_utils.attention_flops_calculator(
    seqlen,
    hidden_size,
    num_attention_heads,
    num_query_groups,
    kv_channels: int | None = None,
    is_swa: bool = False,
    swa_window_size: int = 128
)
```

Calculate the flops for the attention part.

```python
nemo_automodel.components.utils.flops_utils.bert_flops(
    config,
    gbs = 1,
    seq_len = None
)
```

Model FLOPs for BERT family - accepts either AutoConfig or normalized config

```python
nemo_automodel.components.utils.flops_utils.calculate_mfu(
    tflops: float,
    world_size: int,
    time_seconds: float,
    reference_mfu: float | None = None
) -> float
```

Calculate Model FLOPs Utilization (MFU).

**Parameters:**

**`tflops`** `float`

Total TFLOPs across all devices for the measured step.

---

**`world_size`** `int`

Total number of GPUs.

---

**`time_seconds`** `float`

Time taken for computation.

---

**`reference_mfu`** `float | None` — default: None

Peak TFLOPs/s per device for the training precision. The
legacy default is the H100 dense-FP8 peak.

---

**Returns:** `float`

MFU as a percentage.

```python
nemo_automodel.components.utils.flops_utils.clip_vit_l_flops(
    config
)
```

Model FLOPs for CLIP ViT

```python
nemo_automodel.components.utils.flops_utils.deepseek_v41_flops(
    config: typing.Any,
    gbs: int = 1,
    seq_len: int | None = None
) -> float
```

Model FLOPs for DeepSeek-V4.1 (CSA2 sparse attention, single-pass mHC, Engram, MoE).

Accepts `DeepseekV41TextConfig` or the multimodal `DeepseekV41Config` wrapper (its
`text_config` is used; the vision tower is not counted, as for other VL entries). The
module shapes follow `nemo_automodel/components/models/deepseek_v41`:

* attention linears per layer: `wq_a` (hidden x q\_lora\_rank), `wq_b` (q\_lora\_rank x
  heads*head\_dim), `wkv` (hidden x head\_dim, one shared latent), grouped `wo_a`
  (heads*head\_dim -> o\_groups*o\_lora\_rank) and `wo_b` (o\_groups*o\_lora\_rank x hidden);
* compressors on `kv_source_layer_ids`: `wkv` and, for ratio > 1, `wgate` (hidden x
  head\_dim each, applied to every token before pooling);
* sparse attention BMMs: every query attends to `min(i+1, sliding_window)` local keys plus
  `min(index_topk, floor((i+1)/ratio))` selected compressed keys (SWA-only layers have
  `compress_ratios[i] == 0`); QK^T and PV each cost `heads * head_dim` MACs per key;
* the indexer is frozen (`_Indexer.requires_grad_(False)`; hard top-k has no gradient), so
  its projections and its causal scoring over the compressed positions are counted forward
  only (2 x MACs) — the implementation scores the full `[S, S/ratio]` block before
  masking, which is executed work but not model work;
* MoE per layer: router GEMM (hidden x n\_routed\_experts) plus `num_experts_per_tok +
  n_shared_experts` experts of `3 * hidden * moe_intermediate_size`;
* single-pass mHC: two coefficient projections per layer (`hc_mult*(hc_mult+2)` x
  `hc_mult*hidden`); the stream collapse/expand are elementwise and not counted;
* Engram on `engram_layer_ids`: the fused key/value projection
  (`(engram_max_ngram_size-1)*engram_n_heads*engram_head_dim` x `hidden*(hc_mult+1)`);
  the table gather, hashing, norms, RoPE and quantize/dequantize boundaries are not GEMMs;
* the FP32 language-model head (hidden x vocab).

Trainable GEMMs and attention BMMs cost 6 x MACs (forward + backward); activation
recomputation is excluded (that is HFU). For the released Flash configuration the formula
counts 16.09B active GEMM parameters per token (model card: 16B activated during decode)
and 105.6 GFLOPs/token at 4096 tokens, 1.095x the dense-FFN identity `6 x active`.

**Parameters:**

**`config`** `Any`

`DeepseekV41TextConfig`, or a `DeepseekV41Config` wrapper whose `text_config` is used.

---

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

Number of sequences per step.

---

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

Tokens per sequence; defaults to `config.max_position_embeddings`.

---

**Returns:** `float`

Model FLOPs for one training step of `gbs` sequences of `seq_len` tokens.

```python
nemo_automodel.components.utils.flops_utils.deepseekv3_flops(
    config,
    gbs = 1,
    seq_len = None
)
```

Model FLOPs for DeepSeek V3 - accepts either AutoConfig or normalized config

```python
nemo_automodel.components.utils.flops_utils.flux_flops(
    config
)
```

Model FLOPs for FLUX

```python
nemo_automodel.components.utils.flops_utils.get_flops_formula_for_hf_config(
    config: typing.Any
) -> typing.Callable | None
```

Get the appropriate FLOPs formula function for a given HuggingFace config.

**Parameters:**

**`config`** `Any`

HuggingFace model config object

---

**Returns:** `Callable | None`

The appropriate FLOPs formula function, or None for an unregistered

```python
nemo_automodel.components.utils.flops_utils.glm4_moe_flops(
    config,
    gbs = 1,
    seq_len = None
)
```

Estimate FLOPs for GLM4 MoE model configurations.

```python
nemo_automodel.components.utils.flops_utils.gpt3_flops(
    config,
    gbs = 1,
    seq_len = None
)
```

Model FLOPs for GPT3 family - accepts either AutoConfig or normalized config

```python
nemo_automodel.components.utils.flops_utils.gpt_oss_flops(
    config,
    gbs = 1,
    seq_len = None
)
```

Model FLOPs for GPT-OSS

```python
nemo_automodel.components.utils.flops_utils.gpt_oss_flops_calculator(
    gbs,
    num_layers,
    seqlen,
    hidden_size,
    num_attention_heads,
    num_query_groups,
    moe_ffn_hidden_size,
    moe_router_topk,
    vocab_size,
    kv_channels: int | None = None,
    swa_window_size: int = 128,
    window_attn_skip_freq: int | None = 2
)
```

Calculate the flops for the GPT-OSS model

```python
nemo_automodel.components.utils.flops_utils.kimi_k3_flops(
    config: typing.Any,
    gbs: int = 1,
    seq_len: int | None = None
) -> float
```

Model FLOPs for Kimi K3 (hybrid KDA / gated-MLA attention + latent MoE with a SiTU shared expert).

Accepts `KimiK3TextConfig` or the multimodal `KimiK3Config` wrapper (its
`text_config` is used; the vision tower is not counted, as for other VL entries).

The layer pattern is read from the config rather than assumed: 1-based
`linear_attn_config["kda_layers"]` marks KDA layers and every other layer
is MLA (with an extra output-gate projection when `mla_use_output_gate`);
`first_k_dense_replace` / `moe_layer_freq` mark dense vs MoE MLPs, using
the same rule as `KimiDecoderLayer`. Routed experts run in a latent space
of `routed_expert_hidden_size` behind shared down/up projections, and the
shared expert is one SiTU MLP of `moe_intermediate_size * num_shared_experts`.
The router GEMM (`hidden_size x num_experts`, about 0.6% of the total), norms
and the scalar attention-residual projections are omitted, matching the other
MoE formulas in this module.

The router top-k field on this config is `num_experts_per_token` (not the
`num_experts_per_tok` that the generic `transformer_flops` fallback probes),
which is why K3 was previously costed as a dense 93-layer transformer at well
under half of its real FLOPs. For the released checkpoint's configuration
(96 attention heads, dense `intermediate_size` 33792) the formula counts
104.0B active matmul parameters per token and 2.78T total, matching the
model card ("104B activated / 2.8T") and the tech report's Table 1 (104.2B /
2.78T) to within 0.2%, which is inside the accounting difference for the
non-GEMM parameters (router, norms, gate biases). The report's
single MTP layer is not shipped in the checkpoint (`num_nextn_predict_layers`
is 0) and is not modelled by Automodel, so it is not counted.

```python
nemo_automodel.components.utils.flops_utils.llama2_flops(
    config,
    gbs = 1,
    seq_len = None
)
```

Model FLOPs for llama2 family - accepts either AutoConfig or normalized config

```python
nemo_automodel.components.utils.flops_utils.llama3_flops(
    config,
    gbs = 1,
    seq_len = None
)
```

Model FLOPs for llama3 family - accepts either AutoConfig or normalized config

```python
nemo_automodel.components.utils.flops_utils.loss_flops_calculator(
    seqlen,
    hidden_size,
    vocab_size
)
```

Calculate the flops for the loss

```python
nemo_automodel.components.utils.flops_utils.minimax_m2_flops(
    config,
    gbs = 1,
    seq_len = None
)
```

Model FLOPs for MiniMax-M2 family - accepts either AutoConfig or normalized config.

Architecture: GQA attention (Q/K/V/O separate projections, head\_dim may differ from
hidden\_size // num\_heads) + MoE with SwiGLU (no shared experts by default).
Optionally includes MTP (Multi-Token Prediction) modules gated by use\_mtp.

```python
nemo_automodel.components.utils.flops_utils.mixtral_flops(
    config,
    gbs = 1,
    seq_len = None
)
```

Model FLOPs for mixtral family - accepts either AutoConfig or normalized config

```python
nemo_automodel.components.utils.flops_utils.mla_moe_flops(
    config,
    gbs = 1,
    seq_len = None
)
```

Model FLOPs for MLA + MoE models (Kimi K2, GLM-5, Mistral Small 4, etc.).

Handles VL wrappers by extracting text\_config if present.

```python
nemo_automodel.components.utils.flops_utils.moe_mlp_flops_calculator(
    seqlen,
    hidden_size,
    moe_ffn_hidden_size,
    moe_router_topk,
    gated_linear_unit: bool = True
)
```

Calculate the flops for the MLP

```python
nemo_automodel.components.utils.flops_utils.nemotron_flops(
    config,
    gbs = 1,
    seq_len = None
)
```

Model FLOPs for nemotron family - accepts either AutoConfig or normalized config

```python
nemo_automodel.components.utils.flops_utils.nemotronh_flops(
    config,
    gbs = 1,
    seq_len = None
)
```

Model FLOPs for NemotronH

```python
nemo_automodel.components.utils.flops_utils.neva_projection_flops(
    config
)
```

Model FLOPs for NeVA Projection

```python
nemo_automodel.components.utils.flops_utils.qwen3_5_flops(
    config,
    gbs = 1,
    seq_len = None
)
```

Model FLOPs for Qwen3.5 family (MoE and Dense) with hybrid GDN/full attention.

Qwen3.5 uses a hybrid attention pattern: 75% GDN (linear attention) layers
and 25% standard GQA (full attention) layers (full\_attention\_interval=4).
Supports both the MoE variant (Qwen3.5-35B-A3B) and Dense variant (Qwen3.5-27B).

```python
nemo_automodel.components.utils.flops_utils.qwen3_8_flash_next_flops(
    config: typing.Any,
    gbs: int = 1,
    seq_len: int | None = None
) -> float
```

Model FLOPs for Qwen3.8-Flash-Next (GDN + compressed-block QSA, HyperConnections, Engram PLE, MoE).

Accepts `Qwen3_8_FlashNextTextConfig` or the multimodal `Qwen3_8_FlashNextConfig` wrapper
(its `text_config` is used; the vision tower is not loaded by the training path). The module
shapes follow `nemo_automodel/components/models/qwen3_8_flash_next`:

* GatedDeltaNet layers (`layers_block_type == "linear_attention"`): the shared
  `_gdn_attention_per_layer_flops` term (QKV/Z/B/A projections, causal conv, chunked
  delta-rule recurrence, output projection);
* QSA full-attention layers: gated `q_proj` (hidden x 2*heads*head\_dim), `k_proj`/`v_proj`
  (hidden x kv\_heads\*head\_dim) and `o_proj`; the sparse GQA BMMs where the query at zero-based
  position `t` attends to `ratio * min(indexer_budget / ratio, floor((t+1)/ratio))` routed tokens
  plus the `(t+1) mod ratio` tokens of its incomplete causal tail, QK^T and PV each costing
  `heads * head_dim` MACs per key;
* the QSA indexer is frozen (`requires_grad_(False)`; hard top-k has no gradient), so its
  `index_qk_proj` and its causal scoring of `indexer_n_heads` queries against the
  `floor((t+1)/ratio)` compressed keys are counted forward only (2 x MACs);
* MoE per layer: router GEMM (hidden x num\_experts), `num_experts_per_tok` routed plus one shared
  expert of `3 * hidden * moe_intermediate_size` / `shared_expert_intermediate_size`, and the
  shared-expert gate (hidden x 1);
* HyperConnections: two mixers per layer (attention and MoE), each with `input_mix_weight_down`
  (hc\_count*hidden x hc\_lowrank), `input_mix_weight_up` (hc\_lowrank x hc\_count*hidden) and
  `block_inject_weight` (hc\_count\*hidden x hc\_count); the final read mixer has no inject weight;
* Engram PLE on `ple_layer_ids`: `key_proj` (ple\_embed\_dim x hc\_count*hidden), `value_proj`
  (ple\_embed\_dim x hidden) and the depthwise causal convolution (hc\_count*hidden x kernel); the
  table gather, hashing and norms are not GEMMs;
* the untied language-model head (hidden x vocab). The checkpoint's MTP head is not loaded.

Trainable GEMMs and attention BMMs cost 6 x MACs (forward + backward); activation recomputation
is excluded (that is HFU). On the released configuration the formula counts 6.14B active GEMM
parameters per token excluding the LM head (model card: 6B activated), implies a 125.8B backbone
excluding the 51.2B Engram table (model card: 125B), and gives 42.0 GFLOPs/token at 4096 tokens,
1.034x the dense identity `6 x active`.

**Parameters:**

**`config`** `Any`

`Qwen3_8_FlashNextTextConfig`, or a `Qwen3_8_FlashNextConfig` wrapper whose
`text_config` is used.

---

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

Number of sequences per step.

---

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

Tokens per sequence; defaults to `config.max_position_embeddings`.

---

**Returns:** `float`

Model FLOPs for one training step of `gbs` sequences of `seq_len` tokens.

```python
nemo_automodel.components.utils.flops_utils.qwen3_flops(
    config,
    gbs = 1,
    seq_len = None
)
```

Model FLOPs for Qwen3 family - accepts either AutoConfig or normalized config

```python
nemo_automodel.components.utils.flops_utils.step3_5_flash_flops(
    config,
    gbs = 1,
    seq_len = None
)
```

Model FLOPs for Step3.5-Flash (GQA + sliding-window / full attention + MoE).

Architecture: hybrid full/SWA attention with different head counts per type,
MoE with shared expert on most layers, first few layers dense, SwiGLU.

```python
nemo_automodel.components.utils.flops_utils.transformer_flops(
    config,
    gbs = 1,
    seq_len = None
)
```

Calculate FLOPs for a standard Transformer model - accepts either AutoConfig or normalized config.
Note: This does not cover encoder-decoder models.