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

## Module Contents

### Classes

| Name                                                                                                              | Description                                                              |
| ----------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------ |
| [`GroupedExperts`](#nemo_automodel-components-moe-experts-GroupedExperts)                                         | Sparse MoE implementation using all-gather/reduce-scatter primitives.    |
| [`GroupedExpertsDeepEP`](#nemo_automodel-components-moe-experts-GroupedExpertsDeepEP)                             | Sparse MoE implementation using grouped GEMM with DeepEP token dispatch. |
| [`GroupedExpertsTE`](#nemo_automodel-components-moe-experts-GroupedExpertsTE)                                     | MoE experts using TE's GroupedLinear module directly.                    |
| [`_AllGatherConcatVarlenFn`](#nemo_automodel-components-moe-experts-_AllGatherConcatVarlenFn)                     | All-gather with variable local lengths and autograd-safe backward.       |
| [`_ChunkedBiasAdd`](#nemo_automodel-components-moe-experts-_ChunkedBiasAdd)                                       | Add expert bias without materializing a full-size expanded bias.         |
| [`_DeterministicBiasRepeatInterleave`](#nemo_automodel-components-moe-experts-_DeterministicBiasRepeatInterleave) | Expand expert biases while reducing their gradients deterministically.   |

### Functions

| Name                                                                                                                  | Description                                                                                             |
| --------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------- |
| [`_allocate_triton_workspace`](#nemo_automodel-components-moe-experts-_allocate_triton_workspace)                     | Allocate device workspace required by Triton tensor descriptors.                                        |
| [`_apply_bias`](#nemo_automodel-components-moe-experts-_apply_bias)                                                   | Apply per-expert bias to grouped GEMM output.                                                           |
| [`_checkpointed_chunked_expert_mlp`](#nemo_automodel-components-moe-experts-_checkpointed_chunked_expert_mlp)         | Run a grouped expert MLP in bounded row chunks.                                                         |
| [`_deterministic_bias_grad_kernel`](#nemo_automodel-components-moe-experts-_deterministic_bias_grad_kernel)           | Reduce one expert's contiguous token gradients without atomic writes.                                   |
| [`_init_weights`](#nemo_automodel-components-moe-experts-_init_weights)                                               | -                                                                                                       |
| [`_permute_tokens_for_grouped_mm`](#nemo_automodel-components-moe-experts-_permute_tokens_for_grouped_mm)             | Permute tokens by expert assignment and compute offs for torch.\_grouped\_mm.                           |
| [`_reduce_bias_grad_chunked_fallback`](#nemo_automodel-components-moe-experts-_reduce_bias_grad_chunked_fallback)     | Reduce bias gradients without full-size intermediates.                                                  |
| [`_reduce_bias_grad_fallback`](#nemo_automodel-components-moe-experts-_reduce_bias_grad_fallback)                     | Reduce grouped token gradients without requiring Triton.                                                |
| [`_reduce_bias_grad_triton`](#nemo_automodel-components-moe-experts-_reduce_bias_grad_triton)                         | Reduce grouped CUDA token gradients with deterministic compensated FP32 sums.                           |
| [`_resolve_m_splits`](#nemo_automodel-components-moe-experts-_resolve_m_splits)                                       | Resolve the per-expert split sizes for TE's grouped GEMM.                                               |
| [`_segment_reduce_bias_grad`](#nemo_automodel-components-moe-experts-_segment_reduce_bias_grad)                       | Reduce contiguous expert segments into an accumulation-dtype tensor.                                    |
| [`_stabilize_empty_routing_probs_dtype`](#nemo_automodel-components-moe-experts-_stabilize_empty_routing_probs_dtype) | Keep empty expert probabilities stable across checkpoint recomputation.                                 |
| [`_swiglu_clamped_deepep_compiled`](#nemo_automodel-components-moe-experts-_swiglu_clamped_deepep_compiled)           | Clamped SwiGLU (DeepSeek V4 style) for DeepEP.                                                          |
| [`_torch_mm_experts_fwd`](#nemo_automodel-components-moe-experts-_torch_mm_experts_fwd)                               | -                                                                                                       |
| [`_weighted_bias_add_kernel`](#nemo_automodel-components-moe-experts-_weighted_bias_add_kernel)                       | Add weighted expert bias without expanded bias or FP32 activation buffers.                              |
| [`_weighted_bias_probs_grad_kernel`](#nemo_automodel-components-moe-experts-_weighted_bias_probs_grad_kernel)         | Reduce each row's bias-weighted gradient without expanded bias buffers.                                 |
| [`get_expert_activation_for_deepep`](#nemo_automodel-components-moe-experts-get_expert_activation_for_deepep)         | Return the DeepEP expert activation function selected by the MoE config.                                |
| [`is_gated_activation`](#nemo_automodel-components-moe-experts-is_gated_activation)                                   | Check if activation requires gating (gate\_proj + up\_proj).                                            |
| [`quick_geglu_deepep`](#nemo_automodel-components-moe-experts-quick_geglu_deepep)                                     | Apply DeepEP Quick-GEGLU activation and routing probabilities.                                          |
| [`relu2_deepep`](#nemo_automodel-components-moe-experts-relu2_deepep)                                                 | ReLU² activation for DeepEP: relu(x)^2                                                                  |
| [`swiglu_clamped_deepep`](#nemo_automodel-components-moe-experts-swiglu_clamped_deepep)                               | Clamped SwiGLU (DeepSeek V4 style) for DeepEP: Triton kernels on CUDA, the compiled function elsewhere. |
| [`swiglu_oai_deepep`](#nemo_automodel-components-moe-experts-swiglu_oai_deepep)                                       | SwiGLU-OAI (GPT-OSS / MiniMax-M3) activation for grouped experts.                                       |

### Data

[`_BIAS_CHUNK_ROWS`](#nemo_automodel-components-moe-experts-_BIAS_CHUNK_ROWS)

[`_BIAS_CHUNK_THRESHOLD`](#nemo_automodel-components-moe-experts-_BIAS_CHUNK_THRESHOLD)

[`_BIAS_GRAD_TRITON_AVAILABLE`](#nemo_automodel-components-moe-experts-_BIAS_GRAD_TRITON_AVAILABLE)

[`_EXPERT_MLP_CHUNK_BYTES`](#nemo_automodel-components-moe-experts-_EXPERT_MLP_CHUNK_BYTES)

### API

```python
class nemo_automodel.components.moe.experts.GroupedExperts(
    config: nemo_automodel.components.moe.config.MoEConfig,
    backend: typing.Optional[nemo_automodel.components.models.common.utils.BackendConfig] = None
)
```

**Bases:** `Module`

Sparse MoE implementation using all-gather/reduce-scatter primitives.

Supports two compute backends:

* Per-expert loop with gather/scatter (default)
* torch.\_grouped\_mm with argsort-based permutation (backend.experts="torch\_mm";
  "gmm" remains a deprecated compatibility alias)

**`down_proj_bias`**

---

**`down_projs`** `nn.Parameter`

Linear layer for hidden-to-output transformation.

---

**`expert_activation_grouped`** `= get_expert_activation_for_deepep(config)`

---

**`expert_bias`** `= config.expert_bias`

---

**`gate_and_up_projs`** `nn.Parameter`

Linear layer for gate+up (gated) or just up (non-gated).

---

**`gate_up_proj_bias`**

---

**`is_gated`** `= is_gated_activation(config.expert_activation)`

---

**`n_routed_experts`** `int = config.n_routed_experts`

Total number of experts in the model.

---

**`use_mxfp8`**

---

**`use_torch_mm`**

---

```python
nemo_automodel.components.moe.experts.GroupedExperts._forward_grouped_mm(
    x,
    token_mask,
    weights,
    indices,
    gate_and_up_projs,
    down_projs,
    gate_up_proj_bias,
    down_proj_bias,
    n_local_experts,
    experts_start_idx
)
```

Grouped GEMM forward path using torch.\_grouped\_mm.

```python
nemo_automodel.components.moe.experts.GroupedExperts._forward_loop(
    x,
    weights,
    indices,
    token_mask,
    gate_and_up_projs,
    down_projs,
    gate_up_proj_bias,
    down_proj_bias,
    n_local_experts,
    experts_start_idx,
    experts_end_idx
)
```

Per-expert loop forward path using gather/scatter.

```python
nemo_automodel.components.moe.experts.GroupedExperts._gather_ep_inputs(
    x: torch.Tensor,
    token_mask: torch.Tensor,
    weights: torch.Tensor,
    indices: torch.Tensor,
    ep_group: torch.distributed.ProcessGroup,
    ep_size: int
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, list[int]]
```

staticmethod

Gather unequal token shards across the expert-parallel group.

**Parameters:**

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

Tensor of shape \[local\_tokens, hidden], containing local activations.

---

**`token_mask`** `torch.Tensor`

Tensor of shape \[local\_tokens], marking valid tokens.

---

**`weights`** `torch.Tensor`

Tensor of shape \[local\_tokens, topk], containing routing probabilities.

---

**`indices`** `torch.Tensor`

Tensor of shape \[local\_tokens, topk], containing global expert IDs.

---

**`ep_group`** `dist.ProcessGroup`

Expert-parallel process group.

---

**`ep_size`** `int`

Number of ranks in ep\_group.

---

**Returns:** `torch.Tensor`

Activations \[gathered\_tokens, hidden], routing weights \[gathered\_tokens, topk],

```python
nemo_automodel.components.moe.experts.GroupedExperts.forward(
    x: torch.Tensor,
    token_mask: torch.Tensor,
    weights: torch.Tensor,
    indices: torch.Tensor
) -> torch.Tensor
```

Forward pass for the grouped experts.

**Parameters:**

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

Input tensor. Shape is \[num\_tokens, model\_dim].

---

**`token_mask`** `torch.Tensor`

Boolean mask indicating valid tokens.
Shape is \[num\_tokens].

---

**`weights`** `torch.Tensor`

Routing weights for the selected experts.
Shape is \[num\_tokens, num\_activated\_experts].

---

**`indices`** `torch.Tensor`

Indices of the selected experts.
Shape is \[num\_tokens, num\_activated\_experts].

---

**Returns:** `torch.Tensor`

torch.Tensor: Output tensor after expert computation.
Shape is \[num\_tokens, model\_dim]

```python
nemo_automodel.components.moe.experts.GroupedExperts.init_weights(
    buffer_device: torch.device,
    init_std: float = 0.02
) -> None
```

```python
class nemo_automodel.components.moe.experts.GroupedExpertsDeepEP(
    config: nemo_automodel.components.moe.config.MoEConfig,
    backend: typing.Optional[nemo_automodel.components.models.common.utils.BackendConfig] = None,
    dispatcher_backend: str = 'deepep',
    dispatcher_num_sms: int = 20,
    dispatcher_share_token_dispatcher: bool = True,
    dispatcher_async_dispatch: bool = False
)
```

**Bases:** `Module`

Sparse MoE implementation using grouped GEMM with DeepEP token dispatch.

Uses torch.\_grouped\_mm for the canonical `experts="torch_mm"` backend.
`experts="gmm"` remains a deprecated compatibility alias.

Once the experts for a particular token have been identified, this module
is invoked to compute and average the output of the activated experts.

**`dispatcher_hybridep_compact_routing`**

---

**`dispatcher_hybridep_num_blocks_permute`**

---

**`dispatcher_hybridep_num_blocks_unpermute`**

---

**`dispatcher_hybridep_num_sms_preprocessing`**

---

**`dispatcher_hybridep_permute_fusion`**

---

**`down_proj_bias`**

---

**`down_projs`** `nn.Parameter`

Linear layer for hidden-to-output transformation.

---

**`expert_activation`** `= get_expert_activation_for_deepep(config)`

---

**`expert_bias`** `= config.expert_bias`

---

**`gate_and_up_projs`** `nn.Parameter`

Linear layer for gate+up (gated) or just up (non-gated).

---

**`gate_up_proj_bias`**

---

**`is_gated`** `= is_gated_activation(config.expert_activation)`

---

**`n_routed_experts`** `int`

Total number of experts in the model.

---

**`static_routing`**

---

**`use_mxfp8`**

---

```python
nemo_automodel.components.moe.experts.GroupedExpertsDeepEP._dispatch_tokens(
    x: torch.Tensor,
    token_mask: torch.Tensor,
    weights: torch.Tensor,
    indices: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]
```

Dispatch valid tokens and stabilize empty probabilities for recomputation.

**Parameters:**

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

Tensor of shape \[local\_tokens, hidden], containing local activations.

---

**`token_mask`** `torch.Tensor`

Tensor of shape \[local\_tokens], marking valid tokens.

---

**`weights`** `torch.Tensor`

Tensor of shape \[local\_tokens, topk], containing routing probabilities.

---

**`indices`** `torch.Tensor`

Tensor of shape \[local\_tokens, topk], containing global expert IDs.

---

**Returns:** `torch.Tensor`

Expert-sorted activations \[dispatched\_tokens, hidden], counts \[local\_experts],

```python
nemo_automodel.components.moe.experts.GroupedExpertsDeepEP._init_deepep_buffer(
    ep_group: torch.distributed.ProcessGroup
) -> None
```

Initialize DeepEP communication buffers before activation checkpointing.

```python
nemo_automodel.components.moe.experts.GroupedExpertsDeepEP.forward(
    x: torch.Tensor,
    token_mask: torch.Tensor,
    weights: torch.Tensor,
    indices: torch.Tensor
) -> torch.Tensor
```

Forward pass for the grouped experts.

**Parameters:**

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

Input tensor. Shape is \[num\_tokens, model\_dim].

---

**`token_mask`** `torch.Tensor`

Boolean mask indicating valid tokens.
Shape is \[num\_tokens].

---

**`weights`** `torch.Tensor`

Routing weights for the selected experts.
Shape is \[num\_tokens, num\_activated\_experts].

---

**`indices`** `torch.Tensor`

Indices of the selected experts.
Shape is \[num\_tokens, num\_activated\_experts].

---

**Returns:** `torch.Tensor`

torch.Tensor: Output tensor after expert computation.
Shape is \[num\_tokens, model\_dim]

```python
nemo_automodel.components.moe.experts.GroupedExpertsDeepEP.get_pipeline_runtime_initializers() -> tuple[nemo_automodel.components.moe.megatron.token_dispatcher.HybridEPPipelineRuntimeInitializer, ...]
```

Return dispatcher-owned resources that must be ready before a PP step.

**Returns:** `tuple[HybridEPPipelineRuntimeInitializer, ...]`

HybridEP runtime initializers supplied by this expert module's token dispatcher.

```python
nemo_automodel.components.moe.experts.GroupedExpertsDeepEP.init_token_dispatcher(
    ep_mesh: torch.distributed.device_mesh.DeviceMesh
)
```

```python
nemo_automodel.components.moe.experts.GroupedExpertsDeepEP.init_weights(
    buffer_device: torch.device,
    init_std: float = 0.02
) -> None
```

```python
class nemo_automodel.components.moe.experts.GroupedExpertsTE(
    config: nemo_automodel.components.moe.config.MoEConfig,
    backend: typing.Optional[nemo_automodel.components.models.common.utils.BackendConfig] = None,
    dispatcher_backend: str = 'deepep',
    dispatcher_num_sms: int = 20,
    dispatcher_share_token_dispatcher: bool = True,
    dispatcher_async_dispatch: bool = False
)
```

**Bases:** `Module`

MoE experts using TE's GroupedLinear module directly.

Uses TE's native GroupedLinear for computation, providing:

* Optimized grouped GEMM kernels from TE

For expert parallelism, each rank creates GroupedLinear with
num\_local\_experts = n\_routed\_experts / ep\_size.

**`_static_m_splits`** `list | None = None`

---

**`dim`** `= config.dim`

---

**`dispatcher_hybridep_compact_routing`**

---

**`dispatcher_hybridep_num_blocks_permute`**

---

**`dispatcher_hybridep_num_blocks_unpermute`**

---

**`dispatcher_hybridep_num_sms_preprocessing`**

---

**`dispatcher_hybridep_permute_fusion`**

---

**`down_linear`** `GroupedLinear`

Down projection.

---

**`down_proj_bias`** `Tensor | None`

---

**`down_projs`** `Tensor`

---

**`ep_rank`** `= 0`

---

**`expert_activation`** `= get_expert_activation_for_deepep(config)`

---

**`expert_bias`** `= config.expert_bias`

---

**`fp8_padding`** `= Fp8Padding(config.n_routed_experts)`

---

**`fp8_unpadding`** `= Fp8Unpadding(config.n_routed_experts)`

---

**`gate_and_up_projs`** `Tensor`

---

**`gate_up_linear`** `GroupedLinear`

Combined gate and up projection.

---

**`gate_up_proj_bias`** `Tensor | None`

---

**`is_gated`** `= is_gated_activation(config.expert_activation)`

---

**`moe_inter_dim`** `= config.moe_inter_dim`

---

**`n_routed_experts`** `int`

Total number of experts in the model.

---

**`num_local_experts`** `= config.n_routed_experts`

---

**`static_routing`**

---

```python
nemo_automodel.components.moe.experts.GroupedExpertsTE._get_stacked_bias(
    linear: transformer_engine.pytorch.GroupedLinear
) -> torch.Tensor | None
```

```python
nemo_automodel.components.moe.experts.GroupedExpertsTE._get_stacked_weight(
    linear: transformer_engine.pytorch.GroupedLinear,
    transpose: bool = False
) -> torch.Tensor
```

```python
nemo_automodel.components.moe.experts.GroupedExpertsTE._load_from_state_dict(
    state_dict: typing.Dict[str, typing.Any],
    prefix: str,
    local_metadata,
    strict,
    missing_keys,
    unexpected_keys,
    error_msgs
)
```

Load state dict with stacked tensors in DeepEP format.

Converts stacked format to TE GroupedLinear's weight\{i} parameters:

* gate\_and\_up\_projs: \[num\_local\_experts, dim, moe\_inter\_dim \* 2]
* down\_projs: \[num\_local\_experts, moe\_inter\_dim, dim]

```python
nemo_automodel.components.moe.experts.GroupedExpertsTE._normalize_moe_mesh(
    moe_mesh: torch.distributed.device_mesh.DeviceMesh | None
) -> torch.distributed.device_mesh.DeviceMesh | None
```

```python
nemo_automodel.components.moe.experts.GroupedExpertsTE._set_stacked_bias(
    linear: transformer_engine.pytorch.GroupedLinear,
    stacked: torch.Tensor
)
```

```python
nemo_automodel.components.moe.experts.GroupedExpertsTE._set_stacked_weight(
    linear: transformer_engine.pytorch.GroupedLinear,
    stacked: torch.Tensor,
    transpose: bool = False
)
```

```python
nemo_automodel.components.moe.experts.GroupedExpertsTE._to_ep_dtensor(
    tensor: torch.Tensor
) -> torch.Tensor
```

```python
nemo_automodel.components.moe.experts.GroupedExpertsTE.forward(
    x: torch.Tensor,
    token_mask: torch.Tensor,
    weights: torch.Tensor,
    indices: torch.Tensor
) -> torch.Tensor
```

Forward pass using TE's GroupedLinear with native FP8 support.

**Parameters:**

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

\[num\_tokens, model\_dim] input tensor

---

**`token_mask`** `torch.Tensor`

\[num\_tokens] boolean mask for valid tokens

---

**`weights`** `torch.Tensor`

\[num\_tokens, num\_activated\_experts] routing weights

---

**`indices`** `torch.Tensor`

\[num\_tokens, num\_activated\_experts] expert indices

---

**Returns:** `torch.Tensor`

\[num\_tokens, model\_dim] output tensor

```python
nemo_automodel.components.moe.experts.GroupedExpertsTE.get_pipeline_runtime_initializers() -> tuple[nemo_automodel.components.moe.megatron.token_dispatcher.HybridEPPipelineRuntimeInitializer, ...]
```

Return dispatcher-owned resources that must be ready before a PP step.

**Returns:** `tuple[HybridEPPipelineRuntimeInitializer, ...]`

HybridEP runtime initializers supplied by this expert module's token dispatcher.

```python
nemo_automodel.components.moe.experts.GroupedExpertsTE.init_token_dispatcher(
    ep_mesh: torch.distributed.device_mesh.DeviceMesh,
    moe_mesh: torch.distributed.device_mesh.DeviceMesh | None = None
)
```

Initialize the token dispatcher for expert parallelism.

Called by the parallelizer after model initialization.

**Parameters:**

**`ep_mesh`** `DeviceMesh`

Device mesh for expert parallelism.

---

```python
nemo_automodel.components.moe.experts.GroupedExpertsTE.init_weights(
    buffer_device: torch.device,
    init_std: float = 0.02
) -> None
```

Initialize weights using reset\_parameters()

```python
nemo_automodel.components.moe.experts.GroupedExpertsTE.set_moe_mesh(
    moe_mesh: torch.distributed.device_mesh.DeviceMesh | None
) -> None
```

```python
nemo_automodel.components.moe.experts.GroupedExpertsTE.state_dict(
    args = (),
    destination = None,
    prefix = '',
    keep_vars = False,
    kwargs = {}
) -> typing.Dict[str, typing.Any]
```

Return state dict with stacked tensors in DeepEP format.

Converts TE GroupedLinear's weight\{i} parameters to stacked format:

* gate\_and\_up\_projs: \[num\_local\_experts, dim, moe\_inter\_dim \* 2]
* down\_projs: \[num\_local\_experts, moe\_inter\_dim, dim]

When EP is enabled, returns DTensors sharded on dimension 0.

```python
class nemo_automodel.components.moe.experts._AllGatherConcatVarlenFn()
```

**Bases:** `Function`

All-gather with variable local lengths and autograd-safe backward.

Backward uses all-reduce + local narrow instead of reduce-scatter to avoid
monitoredBarrier deadlocks observed with mixed FSDP/EP backward collective ordering.

```python
nemo_automodel.components.moe.experts._AllGatherConcatVarlenFn.backward(
    ctx,
    grad_output: torch.Tensor
)
```

staticmethod

```python
nemo_automodel.components.moe.experts._AllGatherConcatVarlenFn.forward(
    ctx,
    local_tensor: torch.Tensor,
    group: torch.distributed.ProcessGroup,
    gathered_lens: list[int],
    max_len: int
)
```

staticmethod

```python
class nemo_automodel.components.moe.experts._ChunkedBiasAdd()
```

**Bases:** `Function`

Add expert bias without materializing a full-size expanded bias.

```python
nemo_automodel.components.moe.experts._ChunkedBiasAdd.backward(
    ctx,
    grad_output
)
```

staticmethod

Compute value, bias, and probability gradients in bounded memory.

**Parameters:**

**`ctx`**

Autograd context containing bias, token counts, and optional probabilities.

---

**`grad_output`**

Upstream gradient of shape \[tokens, hidden].

---

**Returns:**

Gradients for `value`, `bias`, and `probs`; token counts are non-differentiable.

```python
nemo_automodel.components.moe.experts._ChunkedBiasAdd.forward(
    ctx,
    value,
    bias,
    token_counts,
    probs = None,
    reuse_input = False
)
```

staticmethod

Add optionally weighted expert bias with bounded temporary memory.

**Parameters:**

**`ctx`**

Autograd context used to retain the compact inputs.

---

**`value`**

Grouped expert output of shape \[tokens, hidden].

---

**`bias`**

Per-expert bias of shape \[experts, hidden].

---

**`token_counts`**

Tensor of shape \[experts] containing nonnegative token counts.

---

**`probs`** — default: None

Optional FP32 routing probabilities of shape \[tokens, 1].

---

**`reuse_input`** — default: False

Whether the caller transfers ownership of `value` so its storage can hold the result.

---

**Returns:**

Tensor of shape \[tokens, hidden] and the same dtype as `value`.

```python
class nemo_automodel.components.moe.experts._DeterministicBiasRepeatInterleave()
```

**Bases:** `Function`

Expand expert biases while reducing their gradients deterministically.

```python
nemo_automodel.components.moe.experts._DeterministicBiasRepeatInterleave.backward(
    ctx,
    grad_output
)
```

staticmethod

Reduce expanded bias gradients with one deterministic segmented reduction.

**Parameters:**

**`ctx`**

Autograd context containing the forward token counts.

---

**`grad_output`**

Tensor of shape \[tokens, hidden].

---

**Returns:**

Tuple containing the bias gradient of shape \[experts, hidden] and no gradients for token counts or size.

```python
nemo_automodel.components.moe.experts._DeterministicBiasRepeatInterleave.forward(
    ctx,
    bias,
    token_counts,
    output_size
)
```

staticmethod

Expand each expert bias over its contiguous token group.

**Parameters:**

**`ctx`**

Autograd context used to retain the token counts for backward.

---

**`bias`**

Tensor of shape \[experts, hidden].

---

**`token_counts`**

Tensor of shape \[experts] containing nonnegative token counts.

---

**`output_size`**

Total number of grouped tokens.

---

**Returns:**

Tensor of shape \[tokens, hidden], with each expert row repeated for its tokens.

```python
nemo_automodel.components.moe.experts._allocate_triton_workspace(
    size: int,
    alignment: int,
    stream: int | None
) -> torch.Tensor
```

Allocate device workspace required by Triton tensor descriptors.

**Parameters:**

**`size`** `int`

Workspace size in bytes.

---

**`alignment`** `int`

Required byte alignment. CUDA allocations already provide at least 256-byte alignment.

---

**`stream`** `int | None`

CUDA stream pointer requesting the allocation, or `None` outside a stream.

---

**Returns:** `torch.Tensor`

One-dimensional CUDA byte tensor with `size` elements.

```python
nemo_automodel.components.moe.experts._apply_bias(
    value,
    bias,
    tokens_per_expert,
    permuted_probs = None,
    reuse_input = False
)
```

Apply per-expert bias to grouped GEMM output.

NOTE: torch.\_grouped\_mm accepts a `bias` kwarg in its schema but raises
"RuntimeError: Bias not supported yet" as of PyTorch 2.9.0.
Additionally, down projection bias needs weighting by routing probs
(bias \* permuted\_probs) which native bias support wouldn't handle.

**Parameters:**

**`value`**

Output from grouped GEMM, shape \[total\_tokens, features].

---

**`bias`**

Per-expert bias, shape \[num\_experts, features].

---

**`tokens_per_expert`**

Token counts, shape \[num\_experts].

---

**`permuted_probs`** — default: None

Optional routing probabilities broadcastable to
\[total\_tokens, features], typically \[total\_tokens, 1].

---

**`reuse_input`** — default: False

Reuse `value` storage for the result when the bounded-memory path supports it. The caller
must not read `value` afterward or pass storage retained by a selective-checkpoint cache.

---

**Returns:**

Grouped GEMM output with per-expert bias applied, shape

```python
nemo_automodel.components.moe.experts._checkpointed_chunked_expert_mlp(
    hidden_states: torch.Tensor,
    gate_and_up_projs: torch.Tensor,
    down_projs: torch.Tensor,
    gate_up_proj_bias: torch.Tensor | None,
    down_proj_bias: torch.Tensor | None,
    tokens_per_expert: torch.Tensor,
    permuted_probs: torch.Tensor,
    activation_fn,
    apply_router_weight_after_down: bool
) -> torch.Tensor
```

Run a grouped expert MLP in bounded row chunks.

Each chunk is independently checkpointed, so backward recomputes its two
grouped GEMMs instead of retaining full `[tokens, 2 * inter_dim]` and
`[tokens, inter_dim]` activations. The final `[tokens, hidden]` tensor
is allocated once and populated by slice assignment; no concatenation or
stateful scratch buffer is used.

The checkpoint callback captures the input dtype rather than the routed
tensor so an enclosing activation checkpoint can release that tensor.

This helper is for the plain `torch._grouped_mm` path. MXFP8 uses a
different quantization contract and remains on its existing whole-dispatch
implementation.

**Parameters:**

**`hidden_states`** `torch.Tensor`

Tensor of shape \[tokens, hidden], with rows grouped contiguously by expert.

---

**`gate_and_up_projs`** `torch.Tensor`

Tensor of shape \[experts, hidden, gate\_up], where `gate_up` stores the fused gate and
up projections in the ordering required by `activation_fn`.

---

**`down_projs`** `torch.Tensor`

Tensor of shape \[experts, intermediate, hidden].

---

**`gate_up_proj_bias`** `torch.Tensor | None`

Optional tensor of shape \[experts, gate\_up].

---

**`down_proj_bias`** `torch.Tensor | None`

Optional tensor of shape \[experts, hidden].

---

**`tokens_per_expert`** `torch.Tensor`

Tensor of shape \[experts] containing the number of contiguous rows for each expert.

---

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

Tensor of shape \[tokens, 1] containing each routed token's probability.

---

**`activation_fn`**

Callable mapping tensors of shape \[chunk\_tokens, gate\_up] and \[chunk\_tokens, 1] to a tensor
of shape \[chunk\_tokens, intermediate].

---

**`apply_router_weight_after_down`** `bool`

Whether to apply `permuted_probs` after the down projection instead of in
`activation_fn`.

---

**Returns:** `torch.Tensor`

Tensor of shape \[tokens, hidden] in `hidden_states.dtype` with independent storage.

```python
nemo_automodel.components.moe.experts._deterministic_bias_grad_kernel(
    grad_output_ptr,
    token_offsets_ptr,
    probs_ptr,
    grad_bias_ptr,
    n_tokens,
    hidden: tl.constexpr,
    USE_PROBS: tl.constexpr,
    BLOCK_TOKENS: tl.constexpr,
    BLOCK_HIDDEN: tl.constexpr
)
```

Reduce one expert's contiguous token gradients without atomic writes.

`grad_output_ptr` is a contiguous row-major `[tokens, hidden]` allocation and `probs_ptr` is a
contiguous FP32 `[tokens, 1]` allocation when `USE_PROBS` is true. The caller ensures the gradient row
stride is descriptor-aligned and token offsets fit in int32. Each program exclusively owns one contiguous
FP32 `[expert, hidden block]` tile in `grad_bias_ptr`, so the reduction order is deterministic.

```python
nemo_automodel.components.moe.experts._init_weights(
    module,
    buffer_device: torch.device,
    init_std: float = 0.02
)
```

```python
nemo_automodel.components.moe.experts._permute_tokens_for_grouped_mm(
    indices: torch.Tensor,
    weights: torch.Tensor,
    token_mask: torch.Tensor,
    n_local_experts: int,
    experts_start_idx: int,
    return_slot_ids: bool = False
)
```

Permute tokens by expert assignment and compute offs for torch.\_grouped\_mm.

Takes the raw router outputs and produces sorted token IDs, routing weights,
tokens\_per\_expert counts, and cumulative offsets ready for grouped GEMM.

**Returns:**

Token indices sorted by expert assignment.

```python
nemo_automodel.components.moe.experts._reduce_bias_grad_chunked_fallback(
    grad_output: torch.Tensor,
    token_counts: torch.Tensor,
    probs: torch.Tensor | None,
    output_dtype: torch.dtype
) -> torch.Tensor
```

Reduce bias gradients without full-size intermediates.

**Parameters:**

**`grad_output`** `torch.Tensor`

Tensor of shape \[tokens, hidden], grouped contiguously by expert.

---

**`token_counts`** `torch.Tensor`

Tensor of shape \[experts] containing each contiguous segment length.

---

**`probs`** `torch.Tensor | None`

Optional routing probabilities of shape \[tokens, 1].

---

**`output_dtype`** `torch.dtype`

Result dtype.

---

**Returns:** `torch.Tensor`

Tensor of shape \[experts, hidden] with one weighted gradient sum per expert.

```python
nemo_automodel.components.moe.experts._reduce_bias_grad_fallback(
    grad_output: torch.Tensor,
    token_counts: torch.Tensor
) -> torch.Tensor
```

Reduce grouped token gradients without requiring Triton.

**Parameters:**

**`grad_output`** `torch.Tensor`

Tensor of shape \[tokens, hidden], grouped contiguously by expert.

---

**`token_counts`** `torch.Tensor`

Tensor of shape \[experts] containing each contiguous segment length.

---

**Returns:** `torch.Tensor`

Tensor of shape \[experts, hidden] with one gradient sum per expert.

```python
nemo_automodel.components.moe.experts._reduce_bias_grad_triton(
    grad_output: torch.Tensor,
    token_counts: torch.Tensor,
    probs: torch.Tensor | None = None,
    output_dtype: torch.dtype | None = None
) -> torch.Tensor
```

Reduce grouped CUDA token gradients with deterministic compensated FP32 sums.

**Parameters:**

**`grad_output`** `torch.Tensor`

Contiguous CUDA tensor of shape \[tokens, hidden], grouped by expert. Supported dtypes are
FP16, BF16, and FP32.

---

**`token_counts`** `torch.Tensor`

Contiguous CUDA tensor of shape \[experts] containing each segment length.

---

**`probs`** `torch.Tensor | None` — default: None

Optional FP32 routing probabilities of shape \[tokens, 1].

---

**`output_dtype`** `torch.dtype | None` — default: None

Result dtype. Defaults to `grad_output.dtype`.

---

**Returns:** `torch.Tensor`

CUDA tensor of shape \[experts, hidden] and `output_dtype` when provided,

```python
nemo_automodel.components.moe.experts._resolve_m_splits(
    tokens_per_expert: torch.Tensor | list | tuple,
    static_routing: bool,
    cached_m_splits: list | None
) -> list
```

Resolve the per-expert split sizes for TE's grouped GEMM.

**Parameters:**

**`tokens_per_expert`** `torch.Tensor | list | tuple`

Per-local-expert token counts of shape \[num\_local\_experts]
(CUDA tensor from the flex dispatcher, or an already-materialized sequence).

---

**`static_routing`** `bool`

Benchmark-only static-routing mode (see
`BackendConfig.benchmark_static_routing`) under which the counts are
constant for every microbatch.

---

**`cached_m_splits`** `list | None`

The first microbatch's split list, or None before it is known.

---

**Returns:** `list`

The split sizes as a Python list. Under static routing with a warm cache the

```python
nemo_automodel.components.moe.experts._segment_reduce_bias_grad(
    grad_output: torch.Tensor,
    token_counts: torch.Tensor,
    accumulation_dtype: torch.dtype
) -> torch.Tensor
```

Reduce contiguous expert segments into an accumulation-dtype tensor.

**Parameters:**

**`grad_output`** `torch.Tensor`

Tensor of shape \[tokens, hidden], grouped contiguously by expert.

---

**`token_counts`** `torch.Tensor`

Tensor of shape \[experts] containing each contiguous segment length.

---

**`accumulation_dtype`** `torch.dtype`

Dtype used by the segmented sums.

---

**Returns:** `torch.Tensor`

Tensor of shape \[experts, hidden] in `accumulation_dtype`.

```python
nemo_automodel.components.moe.experts._stabilize_empty_routing_probs_dtype(
    permuted_probs: torch.Tensor,
    compute_dtype: torch.dtype
) -> torch.Tensor
```

Keep empty expert probabilities stable across checkpoint recomputation.

DeepEP may return an empty probability tensor in either the router dtype or
the activation dtype. With non-reentrant activation checkpointing, a later
recomputation can choose the other representation, and PyTorch checkpointing rejects
the resulting saved-tensor metadata change even though both tensors contain
zero elements. Empty probabilities have no numerical contribution, so
canonicalize only that case to the dispatched activation dtype. Non-empty
probabilities retain their original (normally FP32) router precision.

```python
nemo_automodel.components.moe.experts._swiglu_clamped_deepep_compiled(
    x,
    permuted_probs,
    limit: float
)
```

Clamped SwiGLU (DeepSeek V4 style) for DeepEP.

Gate is clamped at `max=limit` and up at `(-limit, +limit)` in FP32
before `silu(gate) * up`; the result is multiplied by the permuted
routing probs and cast back.  Matches the official V4 Expert.forward::

gate = self.w1(x).float()
up   = self.w3(x).float()
if self.swiglu\_limit > 0:
up   = torch.clamp(up,   min=-swiglu\_limit, max=swiglu\_limit)
gate = torch.clamp(gate,                     max=swiglu\_limit)
y = F.silu(gate) \* up

`x` has shape `[..., 2 * inter_dim]` with gate in the first half
and up in the second half (same layout as `weighted_bias_swiglu_impl`).

```python
nemo_automodel.components.moe.experts._torch_mm_experts_fwd(
    hidden_states,
    gate_and_up_projs,
    down_projs,
    tokens_per_expert,
    permuted_probs,
    activation_fn,
    use_mxfp8 = False
)
```

```python
nemo_automodel.components.moe.experts._weighted_bias_add_kernel(
    value_ptr,
    bias_ptr,
    expert_ids_ptr,
    probs_ptr,
    output_ptr,
    numel,
    hidden: tl.constexpr,
    value_stride_0: tl.constexpr,
    value_stride_1: tl.constexpr,
    bias_stride_0: tl.constexpr,
    bias_stride_1: tl.constexpr,
    probs_stride_0: tl.constexpr,
    output_stride_0: tl.constexpr,
    output_stride_1: tl.constexpr,
    BLOCK: tl.constexpr
)
```

Add weighted expert bias without expanded bias or FP32 activation buffers.

**Parameters:**

**`value_ptr`**

Tensor of shape \[tokens, hidden] containing grouped expert outputs.

---

**`bias_ptr`**

Tensor of shape \[experts, hidden] containing per-expert biases.

---

**`expert_ids_ptr`**

Contiguous integer tensor of shape \[tokens] mapping each row to its expert.

---

**`probs_ptr`**

FP32 tensor of shape \[tokens, 1] containing routing weights.

---

**`output_ptr`**

Tensor of shape \[tokens, hidden], optionally aliasing `value_ptr` exactly.

---

**`numel`**

Number of output elements.

---

**`hidden`** `tl.constexpr`

Feature count.

---

**`value_stride_0`** `tl.constexpr`

Value row stride.

---

**`value_stride_1`** `tl.constexpr`

Value feature stride.

---

**`bias_stride_0`** `tl.constexpr`

Bias expert stride.

---

**`bias_stride_1`** `tl.constexpr`

Bias feature stride.

---

**`probs_stride_0`** `tl.constexpr`

Probability row stride.

---

**`output_stride_0`** `tl.constexpr`

Output row stride.

---

**`output_stride_1`** `tl.constexpr`

Output feature stride.

---

**`BLOCK`** `tl.constexpr`

Number of elements processed by each program.

---

```python
nemo_automodel.components.moe.experts._weighted_bias_probs_grad_kernel(
    grad_output_ptr,
    bias_ptr,
    expert_ids_ptr,
    grad_probs_ptr,
    hidden: tl.constexpr,
    grad_stride_0: tl.constexpr,
    grad_stride_1: tl.constexpr,
    bias_stride_0: tl.constexpr,
    bias_stride_1: tl.constexpr,
    probs_stride_0: tl.constexpr,
    BLOCK: tl.constexpr
)
```

Reduce each row's bias-weighted gradient without expanded bias buffers.

**Parameters:**

**`grad_output_ptr`**

Upstream gradient of shape \[tokens, hidden].

---

**`bias_ptr`**

Per-expert bias of shape \[experts, hidden].

---

**`expert_ids_ptr`**

Contiguous integer tensor of shape \[tokens] mapping each row to its expert.

---

**`grad_probs_ptr`**

FP32 probability gradient of shape \[tokens, 1].

---

**`hidden`** `tl.constexpr`

Feature count.

---

**`grad_stride_0`** `tl.constexpr`

Upstream-gradient row stride.

---

**`grad_stride_1`** `tl.constexpr`

Upstream-gradient feature stride.

---

**`bias_stride_0`** `tl.constexpr`

Bias expert stride.

---

**`bias_stride_1`** `tl.constexpr`

Bias feature stride.

---

**`probs_stride_0`** `tl.constexpr`

Probability-gradient row stride.

---

**`BLOCK`** `tl.constexpr`

Power-of-two feature extent processed by each program.

---

```python
nemo_automodel.components.moe.experts.get_expert_activation_for_deepep(
    config: nemo_automodel.components.moe.config.MoEConfig
)
```

Return the DeepEP expert activation function selected by the MoE config.

```python
nemo_automodel.components.moe.experts.is_gated_activation(
    activation: str
) -> bool
```

Check if activation requires gating (gate\_proj + up\_proj).

Gated activations (SwiGLU, Quick-GEGLU) use both gate\_proj and up\_proj,
requiring gate\_and\_up\_projs tensor with shape \[n\_experts, dim, 2\*inter\_dim].

Non-gated activations (ReLU²) only use up\_proj, requiring up\_projs tensor
with shape \[n\_experts, dim, inter\_dim] - 50% memory savings.

```python
nemo_automodel.components.moe.experts.quick_geglu_deepep(
    x,
    permuted_probs,
    alpha: float = 1.702,
    limit: float = 7.0,
    linear_offset: float = 1.0
)
```

Apply DeepEP Quick-GEGLU activation and routing probabilities.

```python
nemo_automodel.components.moe.experts.relu2_deepep(
    x,
    permuted_probs
)
```

ReLU² activation for DeepEP: relu(x)^2

For DeepEP with ReLU², x is the output of the up projection (already computed).
x already has shape \[..., inter\_dim] from efficient up\_proj.

```python
nemo_automodel.components.moe.experts.swiglu_clamped_deepep(
    x,
    permuted_probs,
    limit: float
)
```

Clamped SwiGLU (DeepSeek V4 style) for DeepEP: Triton kernels on CUDA, the compiled function elsewhere.

See `nemo_automodel.components.moe.swiglu_triton` for the kernels and the layout guard.

```python
nemo_automodel.components.moe.experts.swiglu_oai_deepep(
    x,
    permuted_probs,
    alpha: float = 1.702,
    limit: float = 7.0
)
```

SwiGLU-OAI (GPT-OSS / MiniMax-M3) activation for grouped experts.

Computes `gate * sigmoid(alpha * gate) * (up + 1)` in fp32 with gate
clamped `max=limit` and up clamped `+/-limit` (when `limit > 0`).

Unlike `quick_geglu_deepep` (which expects an *interleaved* gate/up
layout, `x[..., ::2]` / `x[..., 1::2]`), this reads the *concatenated*
`[gate | up]` layout produced by `MoESplitExpertsStateDictMixin`
(`torch.cat([gate_t, up_t], dim=-1)`), matching sglang's
`swiglu_no_interleaved_with_alpha_and_limit`.

```python
nemo_automodel.components.moe.experts._BIAS_CHUNK_ROWS = 4096
```

```python
nemo_automodel.components.moe.experts._BIAS_CHUNK_THRESHOLD = 12288
```

```python
nemo_automodel.components.moe.experts._BIAS_GRAD_TRITON_AVAILABLE = _HAVE_TRITON and _HAVE_TRITON_LANGUAGE and hasattr(tl, 'make_tensor_descriptor')
```

```python
nemo_automodel.components.moe.experts._EXPERT_MLP_CHUNK_BYTES = 512 * 1024 * 1024
```