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

## 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. |

### 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)

### 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._RW_CHUNK_ROWS = 8192
```

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