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

MXFP4-resident expert storage and model application for frozen MoE experts.

`GroupedExpertsMXFP4` keeps the frozen routed-expert base weights packed as
fp4-e2m1 + e8m0 block scales (the DeepSeek V4 Flash checkpoint format) and
dequantizes on the fly inside the grouped GEMM, instead of holding them in
bf16. This is the storage win for LoRA / frozen-base training of large MoE
models, where the routed experts dominate parameter memory.

`MXFP4ExpertStorageMixin` shares packed storage and base GEMMs between the
frozen and LoRA experts, with Torch or DeepEP token dispatch.

## Module Contents

### Classes

| Name                                                                                                      | Description                                                                   |
| --------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------- |
| [`GroupedExpertsDeepEPMXFP4`](#nemo_automodel-components-moe-quantized_experts-GroupedExpertsDeepEPMXFP4) | Frozen routed experts with mxfp4-resident base weights under DeepEP dispatch. |
| [`GroupedExpertsMXFP4`](#nemo_automodel-components-moe-quantized_experts-GroupedExpertsMXFP4)             | Frozen routed experts with mxfp4-resident base weights and no adapter.        |
| [`MXFP4ExpertStorageMixin`](#nemo_automodel-components-moe-quantized_experts-MXFP4ExpertStorageMixin)     | Packed-mxfp4 base-weight storage and grouped GEMM for routed experts.         |

### Functions

| Name                                                                                                        | Description                                                               |
| ----------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------- |
| [`_to_local`](#nemo_automodel-components-moe-quantized_experts-_to_local)                                   | Return the local shard of a DTensor, or the tensor unchanged.             |
| [`apply_mxfp4_to_moe_experts`](#nemo_automodel-components-moe-quantized_experts-apply_mxfp4_to_moe_experts) | Apply MXFP4-resident storage to the model's common routed-expert modules. |

### Data

[`logger`](#nemo_automodel-components-moe-quantized_experts-logger)

### API

```python
class nemo_automodel.components.moe.quantized_experts.GroupedExpertsDeepEPMXFP4(
    orig_module: nemo_automodel.components.moe.experts.GroupedExpertsDeepEP
)
```

**Bases:** [MXFP4ExpertStorageMixin](#nemo_automodel-components-moe-quantized_experts-MXFP4ExpertStorageMixin), [GroupedExpertsDeepEP](/nemo-automodel/nemo_automodel/components/moe/experts#nemo_automodel-components-moe-experts-GroupedExpertsDeepEP)

Frozen routed experts with mxfp4-resident base weights under DeepEP dispatch.

Drop-in replacement for `GroupedExpertsDeepEP` when the experts are frozen
(e.g. LoRA on attention only). The DeepEP fused all-to-all token dispatch is reused
unchanged — mxfp4 only changes the two post-dispatch grouped GEMMs, which read the
packed base weights via `MXFP4GroupedMM` instead of bf16 `torch._grouped_mm`.

The DeepEP parent always uses native grouped MM.

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

Forward over mxfp4 base weights with DeepEP dispatch.

Preserves the tensor and EP contract of `GroupedExpertsDeepEP.forward`, replacing the two base
`torch._grouped_mm` calls with `MXFP4GroupedMM` over the packed weights.

```python
class nemo_automodel.components.moe.quantized_experts.GroupedExpertsMXFP4(
    orig_module: nemo_automodel.components.moe.experts.GroupedExperts
)
```

**Bases:** [MXFP4ExpertStorageMixin](#nemo_automodel-components-moe-quantized_experts-MXFP4ExpertStorageMixin), [GroupedExperts](/nemo-automodel/nemo_automodel/components/moe/experts#nemo_automodel-components-moe-experts-GroupedExperts)

Frozen routed experts with mxfp4-resident base weights and no adapter.

Drop-in replacement for `GroupedExperts` when the experts are frozen (e.g.
LoRA training that targets only attention). Forward mirrors
`GroupedExperts._forward_grouped_mm` but reads the packed base weights.

**`use_torch_mm`** `= orig_module.use_torch_mm`

---

```python
nemo_automodel.components.moe.quantized_experts.GroupedExpertsMXFP4._forward_grouped_mm_mxfp4(
    x: torch.Tensor,
    token_mask: torch.Tensor,
    weights: torch.Tensor,
    indices: torch.Tensor,
    n_local_experts: int,
    experts_start_idx: int
) -> torch.Tensor
```

Compute the frozen local experts' contribution before the EP reduction.

**Parameters:**

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

Tensor of shape \[tokens, hidden], gathered across the EP group.

---

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

Boolean tensor of shape \[tokens] selecting valid tokens.

---

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

Tensor of shape \[tokens, top\_k] with differentiable routing probabilities.

---

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

Integer tensor of shape \[tokens, top\_k] with global expert IDs.

---

**`n_local_experts`** `int`

Number of experts on this rank.

---

**`experts_start_idx`** `int`

First global expert ID on this rank.

---

**Returns:** `torch.Tensor`

FP32 tensor of shape \[tokens, hidden], before the EP reduction.

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

Compute packed experts with the tensor and EP contract of GroupedExperts.forward.

```python
class nemo_automodel.components.moe.quantized_experts.MXFP4ExpertStorageMixin()
```

Packed-mxfp4 base-weight storage and grouped GEMM for routed experts.

Mixed into a `GroupedExperts` (or `GroupedExpertsLoRA`) subclass. The base
projections `gate_and_up_projs` / `down_projs` are stored as packed fp4
(int8, two e2m1 nibbles per byte) plus `float8_e8m0fnu` block scales, in
checkpoint orientation `[n_experts, out_dim, in_dim]` so the block scales run
along the contraction dim. Floating-point base parameters are dropped once packed.

Meta weights become packed placeholders for direct checkpoint loading.
Materialized weights are quantized immediately during construction.

**`_MXFP4_BASE_NAMES`** `tuple[str, ...] = ('gate_and_up_projs', 'down_projs')`

---

```python
nemo_automodel.components.moe.quantized_experts.MXFP4ExpertStorageMixin._init_mxfp4_storage() -> None
```

Replace floating-point bases with packed placeholders or quantized weights.

```python
nemo_automodel.components.moe.quantized_experts.MXFP4ExpertStorageMixin._init_packed_placeholders() -> None
```

Register packed meta parameters in checkpoint orientation \[experts, out, in].

```python
nemo_automodel.components.moe.quantized_experts.MXFP4ExpertStorageMixin._mxfp4_base_mm(
    x: torch.Tensor,
    name: str,
    offs: torch.Tensor
) -> torch.Tensor
```

Multiply routed activations by a frozen packed base projection.

**Parameters:**

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

Tensor of shape \[tokens, in\_dim], grouped contiguously by local expert.

---

**`name`** `str`

Base projection parameter name.

---

**`offs`** `torch.Tensor`

Int32 tensor of shape \[local\_experts], holding cumulative token counts.

---

**Returns:** `torch.Tensor`

Tensor of shape \[tokens, out\_dim] with the activation dtype and device.

```python
nemo_automodel.components.moe.quantized_experts.MXFP4ExpertStorageMixin._mxfp4_dequant_expert0(
    name: str,
    dtype: torch.dtype
) -> torch.Tensor
```

Dequantize expert 0 of base weight `name` to compute layout `[in, out]`.

```python
nemo_automodel.components.moe.quantized_experts.MXFP4ExpertStorageMixin._pack_base_weights() -> None
```

Quantize materialized base projections and release their floating-point storage.

```python
nemo_automodel.components.moe.quantized_experts.MXFP4ExpertStorageMixin._register_packed_base_weight(
    name: str,
    packed: torch.Tensor,
    scales: torch.Tensor
) -> None
```

Replace one floating-point base projection with frozen packed parameters.

**Parameters:**

**`name`** `str`

Base projection parameter name.

---

**`packed`** `torch.Tensor`

Int8 tensor of shape \[experts, out\_dim, in\_dim // 2].

---

**`scales`** `torch.Tensor`

E8M0 tensor of shape \[experts, out\_dim, in\_dim // 32].
DTensors must retain the base parameter's mesh and expert-axis placement.

---

```python
nemo_automodel.components.moe.quantized_experts.MXFP4ExpertStorageMixin._validate_mxfp4_config() -> None
```

Reject execution modes that the packed expert computation does not implement.

```python
nemo_automodel.components.moe.quantized_experts._to_local(
    t
)
```

Return the local shard of a DTensor, or the tensor unchanged.

```python
nemo_automodel.components.moe.quantized_experts.apply_mxfp4_to_moe_experts(
    model: torch.nn.Module
) -> torch.nn.Module
```

Apply MXFP4-resident storage to the model's common routed-expert modules.

Call this after LoRA injection and before distributed sharding. LoRA-targeted
experts that are already MXFP4-resident are preserved; remaining plain
`GroupedExperts` and `GroupedExpertsDeepEP` modules are replaced in place.

Meta experts receive packed placeholders for direct checkpoint loading, and
their model state-dict adapter must opt into the MXFP4 expert storage format.
Already materialized expert weights are quantized without changing the adapter's
checkpoint loading mode. Model-specific adapters own checkpoint layouts.

```python
nemo_automodel.components.moe.quantized_experts.logger = logging.getLogger(__name__)
```