> 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.moe.optimized_ops

Memory-optimized MoE elementwise ops extracted from `experts.py`.

Chunked custom-autograd router-weight fp32 multiply: computes the identical
fp32 math in row chunks and saves only low-precision inputs, removing the
full-size fp32 intermediates that otherwise pin \~7 GiB blocks per MoE layer
under activation checkpointing.

With `BackendConfig.compile_router_weight` the same Function runs each pass as
one `torch.compile`d kernel over the whole tensor instead of the eager chunk
loop (see `_compile_router_weight_cores`).

## Module Contents

### Classes

| Name                                                                                                | Description                                                         |
| --------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------- |
| [`_RouterWeightMulFunction`](#nemo_automodel-components-moe-optimized_ops-_RouterWeightMulFunction) | Chunked fp32 router-weight multiply that saves only the raw inputs. |

### Functions

| Name                                                                                                        | Description                                                                        |
| ----------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------- |
| [`_apply_router_weight_fp32`](#nemo_automodel-components-moe-optimized_ops-_apply_router_weight_fp32)       | Apply routing probabilities to expert outputs with fp32 arithmetic.                |
| [`_compile_router_weight_cores`](#nemo_automodel-components-moe-optimized_ops-_compile_router_weight_cores) | Wrap the router-weight multiply cores with `torch.compile`.                        |
| [`_rw_bwd_probs_core`](#nemo_automodel-components-moe-optimized_ops-_rw_bwd_probs_core)                     | Gradient w\.r.t. `probs`: `(grad_out * x).sum(-1)` in fp32, cast to `probs_dtype`. |
| [`_rw_bwd_x_core`](#nemo_automodel-components-moe-optimized_ops-_rw_bwd_x_core)                             | Gradient w\.r.t. `x`: `(grad_out * probs)` in fp32, cast to `x_dtype`.             |
| [`_rw_fwd_core`](#nemo_automodel-components-moe-optimized_ops-_rw_fwd_core)                                 | fp32 router-weight multiply: `(x * probs)` in fp32, cast to `out_dtype`.           |

### Data

[`_RW_CHUNK_ROWS`](#nemo_automodel-components-moe-optimized_ops-_RW_CHUNK_ROWS)

[`_RW_CHUNK_THRESHOLD`](#nemo_automodel-components-moe-optimized_ops-_RW_CHUNK_THRESHOLD)

[`_RW_CORES_COMPILED`](#nemo_automodel-components-moe-optimized_ops-_RW_CORES_COMPILED)

[`_rw_bwd_probs_dispatch`](#nemo_automodel-components-moe-optimized_ops-_rw_bwd_probs_dispatch)

[`_rw_bwd_x_dispatch`](#nemo_automodel-components-moe-optimized_ops-_rw_bwd_x_dispatch)

[`_rw_fwd_dispatch`](#nemo_automodel-components-moe-optimized_ops-_rw_fwd_dispatch)

### API

```python
class nemo_automodel.components.moe.optimized_ops._RouterWeightMulFunction()
```

**Bases:** `Function`

Chunked fp32 router-weight multiply that saves only the raw inputs.

The plain `(x.float() * probs.float()).to(dtype)` lets autograd keep
full-size fp32 \[tokens, hidden] intermediates alive for backward (the
upcast input, the product, and a recompute copy under activation
checkpointing). This Function computes the same fp32 multiply in row
chunks and saves only the low-precision inputs. The forward is
bitwise-identical; the backward matches autograd's fp32 chain
(`grad_x = (g_f32 * probs_f32).to(x.dtype)`,
`grad_probs = (g_f32 * x_f32).sum(-1, keepdim=True)` cast to
`probs.dtype`).

```python
nemo_automodel.components.moe.optimized_ops._RouterWeightMulFunction.backward(
    ctx: typing.Any,
    grad_out: torch.Tensor
) -> tuple[torch.Tensor | None, torch.Tensor | None, None, None]
```

staticmethod

Compute chunked fp32 gradients for the router-weight multiply.

**Parameters:**

**`ctx`** `Any`

Autograd context holding the saved inputs.

---

**`grad_out`** `torch.Tensor`

Upstream gradient of shape \[tokens, hidden].

---

**Returns:** `torch.Tensor | None`

Tuple `(grad_x, grad_probs, None, None)` where `grad_x` has

```python
nemo_automodel.components.moe.optimized_ops._RouterWeightMulFunction.forward(
    ctx: typing.Any,
    x: torch.Tensor,
    probs: torch.Tensor,
    out_dtype: torch.dtype,
    save_x: bool
) -> torch.Tensor
```

staticmethod

Multiply expert outputs by routing probabilities in fp32.

**Parameters:**

**`ctx`** `Any`

Autograd context.

---

**`x`** `torch.Tensor`

Expert outputs of shape \[tokens, hidden].

---

**`probs`** `torch.Tensor`

Routing probabilities of shape \[tokens, 1].

---

**`out_dtype`** `torch.dtype`

Output dtype (the dispatcher's expected activation
dtype, or fp32 for the scatter-add reduction path).

---

**`save_x`** `bool`

Whether backward needs `x`. `x` is only consumed by
the `probs` gradient; when `probs` carries no grad (e.g.
FakeBalancedGate emits constant weights), saving it would pin
a full-size \[tokens, hidden] tensor per MoE layer across the
activation-checkpointing backward window for nothing. Callers
pass `probs.requires_grad`.

---

**Returns:** `torch.Tensor`

Tensor of shape \[tokens, hidden] and dtype `out_dtype`.

```python
nemo_automodel.components.moe.optimized_ops._apply_router_weight_fp32(
    output2: torch.Tensor,
    permuted_probs: torch.Tensor,
    compute_dtype: torch.dtype
) -> torch.Tensor
```

Apply routing probabilities to expert outputs with fp32 arithmetic.

Large 2-D row-aligned inputs go through the chunked custom Function so
autograd does not retain full-size fp32 intermediates; every other shape
keeps the plain eager multiply (bitwise-identical result).

**Parameters:**

**`output2`** `torch.Tensor`

Expert down-projection outputs of shape \[tokens, hidden].

---

**`permuted_probs`** `torch.Tensor`

Routing probabilities broadcastable against
`output2`, typically of shape \[tokens, 1].

---

**`compute_dtype`** `torch.dtype`

Output dtype.

---

**Returns:** `torch.Tensor`

`(output2 * permuted_probs)` computed in fp32 and cast to

```python
nemo_automodel.components.moe.optimized_ops._compile_router_weight_cores() -> None
```

Wrap the router-weight multiply cores with `torch.compile`.

Runs once per process, same lazy pattern as Kimi K3's `compile_situ`: the
compiled functions replace the module-level dispatch targets, so every
expert module shares one compiled kernel per pass and repeated model
construction does not recompile. Fused kernels keep no fp32 intermediates,
so the chunk loop (which only exists to bound eager fp32 transients) is
skipped on this path. Numerics are allclose to eager (same fp32 math, one
rounding to the output dtype), not guaranteed bitwise-identical.

```python
nemo_automodel.components.moe.optimized_ops._rw_bwd_probs_core(
    grad_out: torch.Tensor,
    x: torch.Tensor,
    probs_dtype: torch.dtype
) -> torch.Tensor
```

Gradient w\.r.t. `probs`: `(grad_out * x).sum(-1)` in fp32, cast to `probs_dtype`.

```python
nemo_automodel.components.moe.optimized_ops._rw_bwd_x_core(
    grad_out: torch.Tensor,
    probs: torch.Tensor,
    x_dtype: torch.dtype
) -> torch.Tensor
```

Gradient w\.r.t. `x`: `(grad_out * probs)` in fp32, cast to `x_dtype`.

```python
nemo_automodel.components.moe.optimized_ops._rw_fwd_core(
    x: torch.Tensor,
    probs: torch.Tensor,
    out_dtype: torch.dtype
) -> torch.Tensor
```

fp32 router-weight multiply: `(x * probs)` in fp32, cast to `out_dtype`.

```python
nemo_automodel.components.moe.optimized_ops._RW_CHUNK_ROWS = 8192
```

```python
nemo_automodel.components.moe.optimized_ops._RW_CHUNK_THRESHOLD = 12288
```

```python
nemo_automodel.components.moe.optimized_ops._RW_CORES_COMPILED = False
```

```python
nemo_automodel.components.moe.optimized_ops._rw_bwd_probs_dispatch = _rw_bwd_probs_core
```

```python
nemo_automodel.components.moe.optimized_ops._rw_bwd_x_dispatch = _rw_bwd_x_core
```

```python
nemo_automodel.components.moe.optimized_ops._rw_fwd_dispatch = _rw_fwd_core
```