> This page is for version 0.7.0 · 26.10.
> 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.7.0 · 26.10: https://docs.nvidia.com/nemo/automodel/v0.7/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.megatron.token_dispatcher

## Module Contents

### Classes

| Name                                                                                                                                | Description                                                                             |
| ----------------------------------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------------- |
| [`HybridEPPipelineRuntimeInitializer`](#nemo_automodel-components-moe-megatron-token_dispatcher-HybridEPPipelineRuntimeInitializer) | Prepare HybridEP's process-global buffer before pipeline traffic.                       |
| [`MoEFlexTokenDispatcher`](#nemo_automodel-components-moe-megatron-token_dispatcher-MoEFlexTokenDispatcher)                         | Flex token dispatcher supporting DeepEP, HybridEP, and UCCL-EP backends.                |
| [`TokenDispatcherConfig`](#nemo_automodel-components-moe-megatron-token_dispatcher-TokenDispatcherConfig)                           | Configuration for MoE token dispatch and combine backends.                              |
| [`_DeepepManager`](#nemo_automodel-components-moe-megatron-token_dispatcher-_DeepepManager)                                         | A manager class to handle fused all-to-all communication processes for MoE models using |
| [`_DispatchManager`](#nemo_automodel-components-moe-megatron-token_dispatcher-_DispatchManager)                                     | A manager class to handle dispatch and combine processes for MoE models.                |
| [`_HybridEPManager`](#nemo_automodel-components-moe-megatron-token_dispatcher-_HybridEPManager)                                     | A manager class to handle fused all-to-all communication processes for MoE models using |
| [`_HybridEPMetadataProcessor`](#nemo_automodel-components-moe-megatron-token_dispatcher-_HybridEPMetadataProcessor)                 | Graphable conversion of top-k metadata for dense or compact HybridEP dispatch.          |

### Functions

| Name                                                                                                                    | Description                                                                           |
| ----------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------- |
| [`_align_up`](#nemo_automodel-components-moe-megatron-token_dispatcher-_align_up)                                       | Round `count` up to the next multiple of `multiple`.                                  |
| [`_assert_no_hybridep_overflow`](#nemo_automodel-components-moe-megatron-token_dispatcher-_assert_no_hybridep_overflow) | Device-side guard: HybridEP truncates silently when the capacity is exceeded and sets |

### Data

[`_HYBRIDEP_TOKEN_ALIGNMENT`](#nemo_automodel-components-moe-megatron-token_dispatcher-_HYBRIDEP_TOKEN_ALIGNMENT)

[`_STATIC_ROUTING_PAD_PIN`](#nemo_automodel-components-moe-megatron-token_dispatcher-_STATIC_ROUTING_PAD_PIN)

### API

```python
class nemo_automodel.components.moe.megatron.token_dispatcher.HybridEPPipelineRuntimeInitializer(
    manager: nemo_automodel.components.moe.megatron.token_dispatcher._HybridEPManager,
    hidden_dim: int,
    dtype: torch.dtype
)
```

Prepare HybridEP's process-global buffer before pipeline traffic.

**`resource_key`** `tuple[type[HybridEPPipelineRuntimeInitializer], str]`

Identify the single HybridEP buffer owned by this process.

---

**`signature`** `tuple[object, ...]`

Return the immutable configuration served by the global buffer.

---

```python
nemo_automodel.components.moe.megatron.token_dispatcher.HybridEPPipelineRuntimeInitializer._resolve_group_capacity(
    num_tokens: int,
    device: torch.device
) -> int
```

Return the EP-group maximum JIT capacity for this batch.

```python
nemo_automodel.components.moe.megatron.token_dispatcher.HybridEPPipelineRuntimeInitializer.prepare(
    num_tokens: int,
    device: torch.device
) -> None
```

Initialize the group capacity when the upcoming pipeline shape grows.

**Parameters:**

**`num_tokens`** `int`

Maximum number of tokens in the upcoming pipeline microbatch.

---

**`device`** `torch.device`

Runtime device used for group capacity exchange and initialization.

---

**Raises:**

* `ValueError`: If the token capacity is not positive.
* `RuntimeError`: If the process-global resource signature conflicts or initialization fails.

```python
nemo_automodel.components.moe.megatron.token_dispatcher.HybridEPPipelineRuntimeInitializer.required_capacity(
    num_tokens: int
) -> int
```

staticmethod

Match HybridEP's token-capacity floor and rounding for its JIT key.

**Parameters:**

**`num_tokens`** `int`

Requested local token count.

---

**Returns:** `int`

Normalized token capacity used by HybridEP's compiled runtime.

```python
class nemo_automodel.components.moe.megatron.token_dispatcher.MoEFlexTokenDispatcher(
    num_local_experts: int,
    local_expert_indices: typing.List[int],
    config: nemo_automodel.components.moe.megatron.token_dispatcher.TokenDispatcherConfig,
    ep_group: torch.distributed.ProcessGroup
)
```

Flex token dispatcher supporting DeepEP, HybridEP, and UCCL-EP backends.

**`_comm_manager`** `= MoEFlexTokenDispatcher.shared_uccl_manager`

---

**`ep_size`** `= ep_group.size()`

---

**`hybridep_metadata_processor`** `_HybridEPMetadataProcessor | None = None`

---

**`shared_deepep_manager`** `_DeepepManager = None`

---

**`shared_hybridep_manager`** `_HybridEPManager = None`

---

**`shared_uccl_manager`** `_DeepepManager = None`

---

**`tp_size`** `= 1`

---

```python
nemo_automodel.components.moe.megatron.token_dispatcher.MoEFlexTokenDispatcher._initialize_metadata(
    num_local_tokens: int,
    probs: torch.Tensor
) -> torch.Tensor
```

Initialize the routing map and probs to a unified format covering the TPxEP group.
This design decouples the communication group from underlying model parallelism groups,
such that the communication strategy of tokens can be agnostic of TP size and EP size.

```python
nemo_automodel.components.moe.megatron.token_dispatcher.MoEFlexTokenDispatcher.combine_all_to_all(
    hidden_states: torch.Tensor,
    async_finish: bool = True,
    allocate_on_comm_stream: bool = True
)
```

Performs all-to-all communication to combine tokens after expert processing.

```python
nemo_automodel.components.moe.megatron.token_dispatcher.MoEFlexTokenDispatcher.combine_postprocess(
    hidden_states: torch.Tensor
)
```

Post-processes the combined hidden states after all-to-all communication.

This method reshapes the combined hidden states to match the original input shape.

```python
nemo_automodel.components.moe.megatron.token_dispatcher.MoEFlexTokenDispatcher.combine_preprocess(
    hidden_states: torch.Tensor
)
```

Pre-processes the hidden states before combining them after expert processing.

This method restores the hidden states to their original ordering before expert processing
by using the communication manager's restoration function.

```python
nemo_automodel.components.moe.megatron.token_dispatcher.MoEFlexTokenDispatcher.dispatch_all_to_all(
    hidden_states: torch.Tensor,
    probs: torch.Tensor = None,
    async_finish: bool = True,
    allocate_on_comm_stream: bool = True
)
```

Performs all-to-all communication to dispatch tokens across expert parallel ranks.

```python
nemo_automodel.components.moe.megatron.token_dispatcher.MoEFlexTokenDispatcher.dispatch_postprocess(
    hidden_states: torch.Tensor
)
```

Post-processes the dispatched hidden states after all-to-all communication.

This method retrieves the permuted hidden states by experts, calculates the number of tokens
per expert, and returns the processed data ready for expert processing.

```python
nemo_automodel.components.moe.megatron.token_dispatcher.MoEFlexTokenDispatcher.dispatch_preprocess(
    hidden_states: torch.Tensor,
    num_local_tokens: int,
    probs: torch.Tensor
)
```

Preprocesses the hidden states and routing information before dispatching tokens to experts.
Args:
hidden_states (torch.Tensor): Input hidden states to be processed
num_local_tokens (int): Number of tokens to be processed
probs (torch.Tensor): Routing probabilities for each token-expert pair

**Returns:**

Tuple containing:

```python
nemo_automodel.components.moe.megatron.token_dispatcher.MoEFlexTokenDispatcher.dispatch_preprocess2(
    hidden_states: torch.Tensor,
    num_local_tokens: int,
    token_probs: torch.Tensor,
    token_indices: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]
```

Prepare hidden states and routing metadata for dispatch.

DeepEP consumes `token_indices[T, K]` and `token_probs[T, K]` directly. HybridEP
uses its established dense routing map by default. Compact routing can be enabled to
keep top-k indices for metadata communication while expanding probabilities to `[T, E]`
so the existing dense router-gradient layout is preserved.

**Parameters:**

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

Hidden states with shape \[..., hidden].

---

**`num_local_tokens`** `int`

Number of tokens before expert-parallel dispatch.

---

**`token_probs`** `torch.Tensor`

Top-k routing probabilities with shape \[tokens, top_k].

---

**`token_indices`** `torch.Tensor`

Expert indices with shape \[tokens, top_k].

---

**Returns:** `torch.Tensor`

Flattened hidden states with shape \[tokens, hidden] and routing probabilities with

```python
nemo_automodel.components.moe.megatron.token_dispatcher.MoEFlexTokenDispatcher.get_pipeline_runtime_initializers(
    hidden_dim: int,
    dtype: torch.dtype
) -> tuple[nemo_automodel.components.moe.megatron.token_dispatcher.HybridEPPipelineRuntimeInitializer, ...]
```

Return model-owned initializers needed before a pipeline step.

**Parameters:**

**`hidden_dim`** `int`

Hidden dimension dispatched to experts.

---

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

Hidden-state data type used by the expert module.

---

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

A HybridEP initializer for HybridEP dispatchers, otherwise an empty tuple.

```python
nemo_automodel.components.moe.megatron.token_dispatcher.MoEFlexTokenDispatcher.token_permutation(
    hidden_states: torch.Tensor,
    num_local_tokens: int,
    probs: torch.Tensor
) -> typing.Tuple[torch.Tensor, torch.Tensor, torch.Tensor]
```

Permutes tokens according to probs and dispatches them to experts.

This method implements the token permutation process in three steps:

1. Preprocess the hidden states
2. Perform all-to-all communication to dispatch tokens
3. Post-process the dispatched tokens for expert processing

```python
nemo_automodel.components.moe.megatron.token_dispatcher.MoEFlexTokenDispatcher.token_permutation2(
    hidden_states: torch.Tensor,
    num_local_tokens: int,
    token_probs: torch.Tensor,
    token_indices: torch.Tensor
) -> typing.Tuple[torch.Tensor, torch.Tensor, torch.Tensor]
```

Permutes tokens according to probs and dispatches them to experts.

This method implements the token permutation process in three steps:

1. Preprocess the hidden states
2. Perform all-to-all communication to dispatch tokens
3. Post-process the dispatched tokens for expert processing

```python
nemo_automodel.components.moe.megatron.token_dispatcher.MoEFlexTokenDispatcher.token_unpermutation(
    hidden_states: torch.Tensor
) -> torch.Tensor
```

Reverses the token permutation process to restore the original token order.

This method implements the token unpermutation process in three steps:

1. Pre-process the hidden states to restore their original ordering
2. Perform all-to-all communication to combine tokens
3. Post-process the combined tokens to match the original input shape

```python
class nemo_automodel.components.moe.megatron.token_dispatcher.TokenDispatcherConfig(
    moe_enable_deepep: bool = True,
    moe_permute_fusion: bool = False,
    moe_expert_capacity_factor: float | None = None,
    moe_router_topk: int = 2,
    moe_router_expert_pad_multiple: int | None = None,
    num_moe_experts: int = 64,
    moe_router_dtype: str = 'fp32',
    moe_flex_dispatcher_backend: typing.Literal['deepep', 'hybridep', 'uccl_ep'] = 'deepep',
    moe_hybridep_capacity_factor: float | None = None,
    moe_hybridep_equal_token_counts: bool = False,
    moe_deepep_num_sms: int = 20,
    moe_hybridep_num_sms: int = 24,
    moe_hybridep_permute_fusion: bool = False,
    moe_hybridep_compact_routing: bool = False,
    moe_hybridep_num_sms_preprocessing: int | None = None,
    moe_hybridep_num_blocks_permute: int | None = None,
    moe_hybridep_num_blocks_unpermute: int | None = None,
    moe_share_token_dispatcher: bool = True,
    moe_deepep_async_dispatch: bool = False,
    moe_benchmark_static_routing: bool = False
)
```

Dataclass

Configuration for MoE token dispatch and combine backends.

**`moe_benchmark_static_routing`** `bool = False`

Benchmark-only (mirrors BackendConfig.benchmark_static_routing, validated there):
routing is forced-balanced with no noise, so every dispatch permutes the same token
count. Persist num_permuted_tokens across dispatches to keep HybridEP on its
non-blocking size path instead of waiting on a device-side count per microbatch.

---

**`moe_deepep_async_dispatch`** `bool = False`

Use asynchronous DeepEP/UCCL-EP dispatch/combine and communication-stream allocations.

---

**`moe_deepep_num_sms`** `int = 20`

Number of SMs to use for DeepEP backend.

---

**`moe_enable_deepep`** `bool = True`

Enable DeepEP for efficient token dispatching and combine in MoE models.

---

**`moe_expert_capacity_factor`** `float | None = None`

moe_expert_capacity_factor (float): The capacity factor for each expert, None means no token
will be dropped. The default is None.

---

**`moe_flex_dispatcher_backend`** `Literal['deepep', 'hybridep', 'uccl_ep'] = 'deepep'`

Backend for the flex token dispatcher. Options: 'deepep', 'hybridep', or 'uccl_ep'.

---

**`moe_hybridep_capacity_factor`** `float | None = None`

HybridEP capacity mode (mirrors BackendConfig.dispatcher_capacity_factor): after one blocking
calibration dispatch, size every later dispatch to the calibrated permuted row count times this
factor and run it non-blocking. None keeps the blocking path.

---

**`moe_hybridep_compact_routing`** `bool = False`

Send compact top-k indices instead of a dense Boolean HybridEP routing map.

---

**`moe_hybridep_equal_token_counts`** `bool = False`

HybridEP (mirrors BackendConfig.dispatcher_equal_token_counts): every EP rank dispatches the
same row count, so the per-dispatch pad-size all-reduce and its host sync are skipped.

---

**`moe_hybridep_num_blocks_permute`** `int | None = None`

Optional number of HybridEP dispatch permutation blocks.

---

**`moe_hybridep_num_blocks_unpermute`** `int | None = None`

Optional number of HybridEP combine unpermutation blocks.

---

**`moe_hybridep_num_sms`** `int = 24`

Number of SMs to use for HybridEP dispatch and combine APIs.

---

**`moe_hybridep_num_sms_preprocessing`** `int | None = None`

Optional number of SMs used by HybridEP routing-metadata preprocessing.

---

**`moe_hybridep_permute_fusion`** `bool = False`

Fuse HybridEP permutation with dispatch and unpermutation with combine.

---

**`moe_permute_fusion`** `bool = False`

Fuse token rearrangement ops during token dispatching.

---

**`moe_router_dtype`** `str = 'fp32'`

Data type for routing and expert output weighted averaging. Using fp32 or fp64 can
improve stability especially when the number of experts is large (e.g. finegrained-moe).
None means no changes for dtype.

---

**`moe_router_expert_pad_multiple`** `int | None = None`

Number of tokens to pad to a multiple of for each expert.

---

**`moe_router_topk`** `int = 2`

Number of experts to route to for each token.

---

**`moe_share_token_dispatcher`** `bool = True`

Share one communication manager instance across MoE layers for the configured backend.

---

**`num_moe_experts`** `int = 64`

Number of experts to use for MoE layer. When set, it replaces MLP with MoE layer. Set to None
for no MoE.

---

```python
class nemo_automodel.components.moe.megatron.token_dispatcher._DeepepManager(
    group: torch.distributed.ProcessGroup,
    router_topk: int,
    permute_fusion: bool = False,
    capacity_factor: float | None = None,
    num_experts: int | None = None,
    num_local_experts: int | None = None,
    router_dtype: str | None = None,
    moe_router_expert_pad_multiple: int | None = None,
    _dispatch_fn = None,
    _combine_fn = None
)
```

**Bases:** [\_DispatchManager](#nemo_automodel-components-moe-megatron-token_dispatcher-_DispatchManager)

A manager class to handle fused all-to-all communication processes for MoE models using
DeepEP backend. See [https://github.com/deepseek-ai/deepep](https://github.com/deepseek-ai/deepep) for more details.

The workflow of the DeepEP dispatcher is:
(1) setup_metadata(): Process routing map and probabilities to prepare dispatch metadata
(2) dispatch():

* Use fused kernel to permute tokens and perform all-to-all communication in single step
  (3) get_permuted_hidden_states_by_instances():
* Convert routing map and probabilities to multihot format
* Permute tokens using fused kernel
  (4) get_restored_hidden_states_by_instances():
* Reverse permutation using fused kernel
  (5) combine():
* Reverse process using fused kernel to unpermute and perform all-to-all in single step

This implementation uses fused communication kernels (fused_dispatch/fused_combine) that
combine permutation and communication operations for improved efficiency compared to
separate permute+alltoall steps.

**`_fused_combine`**

---

**`_fused_dispatch`**

---

**`dispatched_indices`** `Tensor | None`

---

**`dispatched_probs`** `Tensor | None`

---

**`dispatched_routing_map`** `Tensor | None`

---

**`handle`** `object | None`

---

**`hidden_shape_before_permute`** `Size | None`

---

**`reversed_mapping_for_combine`** `Tensor | None`

---

**`token_indices`** `Tensor | None`

---

**`token_probs`** `Tensor | None`

---

**`tokens_per_expert`** `Tensor | None`

---

```python
nemo_automodel.components.moe.megatron.token_dispatcher._DeepepManager._indices_to_multihot(
    indices,
    probs
)
```

Converts a tensor of indices to a multihot vector.

**Parameters:**

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

\[num_tokens, topk] token indices, where -1 means masked out.

---

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

\[num_tokens, topk] token probabilities.

---

**Returns:**

Tuple\[torch.Tensor, torch.Tensor]:

* routing_map: Multihot vector.
* probs: Multihot probabilities.

```python
nemo_automodel.components.moe.megatron.token_dispatcher._DeepepManager._release_dispatch_state() -> None
```

Initialize or release state owned by one dispatch.

```python
nemo_automodel.components.moe.megatron.token_dispatcher._DeepepManager.combine(
    hidden_states: torch.Tensor,
    async_finish: bool = False,
    allocate_on_comm_stream: bool = False
) -> torch.Tensor
```

Reverse process using fused kernel to unpermute and perform all-to-all in single step

```python
nemo_automodel.components.moe.megatron.token_dispatcher._DeepepManager.dispatch(
    hidden_states: torch.Tensor,
    async_finish: bool = False,
    allocate_on_comm_stream: bool = False
) -> torch.Tensor
```

Dispatch the hidden_states

```python
nemo_automodel.components.moe.megatron.token_dispatcher._DeepepManager.get_dispatched_metadata() -> torch.Tensor
```

```python
nemo_automodel.components.moe.megatron.token_dispatcher._DeepepManager.get_number_of_tokens_per_expert() -> torch.Tensor
```

Get the number of tokens per expert.

```python
nemo_automodel.components.moe.megatron.token_dispatcher._DeepepManager.get_permuted_hidden_states_by_experts(
    hidden_states: torch.Tensor
) -> torch.Tensor
```

* Convert routing map and probabilities to multihot format
* Permute tokens using fused kernel

```python
nemo_automodel.components.moe.megatron.token_dispatcher._DeepepManager.get_restored_hidden_states_by_experts(
    hidden_states: torch.Tensor
) -> torch.Tensor
```

Restore the hidden states to their original ordering before expert processing

```python
nemo_automodel.components.moe.megatron.token_dispatcher._DeepepManager.setup_metadata(
    num_local_tokens: int,
    probs: torch.Tensor
)
```

Process routing map and probabilities to prepare dispatch metadata

```python
class nemo_automodel.components.moe.megatron.token_dispatcher._DispatchManager()
```

Abstract

A manager class to handle dispatch and combine processes for MoE models.

DispatcherManager handles token dispatching according to the routing_map of format
\[num_local_tokens, world_size, num_instances]. The routing_map is a 3D tensor where each
element indicates whether a token should be sent to a specific rank.

num_instances is the maximum number of tokens instances dispatched into a target rank, it
can be the number of local experts, or the size of sub_group.

```python
nemo_automodel.components.moe.megatron.token_dispatcher._DispatchManager._release_dispatch_state() -> None
```

abstract

Initialize or release this manager's per-dispatch state.

Routing probabilities can retain recomputed activations through their
autograd graph. Dispatch and permutation must also release consumed
tensors early because checkpoint recomputation can stop before combine.
Persistent capacity and static-routing state must remain untouched.

```python
nemo_automodel.components.moe.megatron.token_dispatcher._DispatchManager.combine(
    hidden_states: torch.Tensor
) -> torch.Tensor
```

abstract

Combine the hidden_states after expert processing.

```python
nemo_automodel.components.moe.megatron.token_dispatcher._DispatchManager.dispatch(
    hidden_states: torch.Tensor
) -> torch.Tensor
```

abstract

Dispatch the hidden_states according to the routing_map.

```python
nemo_automodel.components.moe.megatron.token_dispatcher._DispatchManager.get_dispatched_metadata() -> torch.Tensor
```

abstract

Get the metadata of the dispatched hidden_states.

```python
nemo_automodel.components.moe.megatron.token_dispatcher._DispatchManager.get_permuted_hidden_states_by_experts(
    hidden_states: torch.Tensor
) -> torch.Tensor
```

abstract

Get the permuted hidden states by instances.

```python
nemo_automodel.components.moe.megatron.token_dispatcher._DispatchManager.get_restored_hidden_states_by_experts(
    hidden_states: torch.Tensor
) -> torch.Tensor
```

abstract

Get the restored hidden states by instances.

```python
nemo_automodel.components.moe.megatron.token_dispatcher._DispatchManager.setup_metadata(
    routing_map: torch.Tensor,
    probs: torch.Tensor
)
```

abstract

Set up metadata of routing_map and probs.

```python
class nemo_automodel.components.moe.megatron.token_dispatcher._HybridEPManager(
    group: torch.distributed.ProcessGroup,
    num_local_experts: int,
    num_experts: int,
    router_topk: int,
    permute_fusion: bool = False,
    moe_hybridep_num_sms: int = 24,
    moe_hybridep_permute_fusion: bool = False,
    moe_hybridep_compact_routing: bool = False,
    moe_hybridep_num_sms_preprocessing: int | None = None,
    moe_hybridep_num_blocks_permute: int | None = None,
    moe_hybridep_num_blocks_unpermute: int | None = None,
    benchmark_static_routing: bool = False,
    moe_hybridep_capacity_factor: float | None = None,
    moe_hybridep_equal_token_counts: bool = False
)
```

**Bases:** [\_DispatchManager](#nemo_automodel-components-moe-megatron-token_dispatcher-_DispatchManager)

A manager class to handle fused all-to-all communication processes for MoE models using
HybridEP backend. See [https://github.com/deepseek-ai/DeepEP/tree/hybrid-ep](https://github.com/deepseek-ai/DeepEP/tree/hybrid-ep) for more details.

The workflow of the HybridEP dispatcher is:
(1) setup_metadata(): Process routing map and probabilities to prepare dispatch metadata
(2) dispatch():

* Permute tokens for communication, perform all-to-all communication,
  and permute tokens for experts in single step
  (3) combine():
* Unpermute tokens for communication, perform all-to-all communication,
  and unpermute tokens for attention in single step

**`_hybridep_capacity`** `int | None = None`

---

**`_static_target_tokens`** `int | None = None`

---

**`dispatched_probs`** `Tensor | None`

---

**`handle`** `object | None`

---

**`num_permuted_tokens`** `int | Tensor | None`

---

**`num_unpadded_tokens`** `int | None`

---

**`routing_map`** `Tensor | None`

---

**`token_indices`** `Tensor | None`

---

**`token_probs`** `Tensor | None`

---

**`tokens_per_expert`** `Tensor | None`

---

```python
nemo_automodel.components.moe.megatron.token_dispatcher._HybridEPManager._calibrate_hybridep_capacity(
    tokens_per_expert: torch.Tensor,
    device: torch.device
) -> None
```

Turn the one blocking calibration dispatch into the capacity every later dispatch uses.

The blocking dispatch already brought `tokens_per_expert` to the host, so reading its sum is
free; the capacity is that count times `hybridep_capacity_factor`, rounded up to a multiple of
`pad_multiple` when one is set (HybridEP's only requirement on the extent in non-blocking mode)
and otherwise to the 4-token alignment as slack, taken as the EP-group maximum once so every rank
allocates the same buffers.

**Parameters:**

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

`[num_local_experts]` integer tensor of routed rows per local expert, as
returned by the blocking calibration dispatch; its sum is read on the host here.

---

**`device`** `torch.device`

Device for the EP-group max all-reduce of the capacity.

---

```python
nemo_automodel.components.moe.megatron.token_dispatcher._HybridEPManager._indices_to_multihot(
    indices: torch.Tensor,
    probs: torch.Tensor
)
```

Converts a tensor of indices to a multihot vector.

```python
nemo_automodel.components.moe.megatron.token_dispatcher._HybridEPManager._release_dispatch_state() -> None
```

Initialize or release state owned by one dispatch.

```python
nemo_automodel.components.moe.megatron.token_dispatcher._HybridEPManager._resolve_num_permuted_tokens(
    capacity_mode: bool
)
```

The `num_permuted_tokens` this dispatch hands to HybridEP, decided in one place.

`None` is the blocking path (HybridEP reads the count on the device and drains the compute
stream); a host value keeps the dispatch non-blocking. The static-routing pin is the first
dispatch's count (a 0-d tensor); capacity mode passes the calibrated capacity (a host `int`)
once calibration has happened. Nothing is cached across dispatches otherwise: a stale value
would be reused under non-reentrant activation checkpointing.

```python
nemo_automodel.components.moe.megatron.token_dispatcher._HybridEPManager.combine(
    hidden_states: torch.Tensor,
    async_finish: bool = True,
    allocate_on_comm_stream: bool = True
) -> torch.Tensor
```

```python
nemo_automodel.components.moe.megatron.token_dispatcher._HybridEPManager.dispatch(
    hidden_states: torch.Tensor,
    async_finish: bool = True,
    allocate_on_comm_stream: bool = True,
    capacity_mode: bool | None = None
) -> torch.Tensor
```

Dispatch hidden states using compact top-k metadata when available.

**Parameters:**

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

Hidden states with shape \[tokens, hidden].

---

**`async_finish`** `bool` — default: True

Ignored because HybridEP does not expose asynchronous completion here.

---

**`allocate_on_comm_stream`** `bool` — default: True

Ignored because HybridEP owns its communication stream.

---

**`capacity_mode`** `bool | None` — default: None

`None` (default) derives it from `hybridep_capacity_factor`; `False`
forces one blocking dispatch that neither uses nor calibrates the capacity (the
runtime warmup).

---

**Returns:** `torch.Tensor`

Dispatched hidden states with shape \[dispatched_tokens, hidden].

```python
nemo_automodel.components.moe.megatron.token_dispatcher._HybridEPManager.get_dispatched_metadata() -> torch.Tensor
```

```python
nemo_automodel.components.moe.megatron.token_dispatcher._HybridEPManager.get_number_of_tokens_per_expert() -> torch.Tensor
```

```python
nemo_automodel.components.moe.megatron.token_dispatcher._HybridEPManager.get_permuted_hidden_states_by_experts(
    hidden_states: torch.Tensor
) -> torch.Tensor
```

```python
nemo_automodel.components.moe.megatron.token_dispatcher._HybridEPManager.get_restored_hidden_states_by_experts(
    hidden_states: torch.Tensor
) -> torch.Tensor
```

```python
nemo_automodel.components.moe.megatron.token_dispatcher._HybridEPManager.initialize_runtime(
    num_tokens: int,
    hidden_dim: int,
    dtype: torch.dtype,
    device: torch.device
) -> None
```

Initialize HybridEP's dispatch/combine kernels for a capacity without running a model stage.

HybridEP JIT-compiles dispatch and combine kernels on first use.  Under
pipeline parallelism, allowing that work to happen in the real schedule
serializes the compile cost across PP stages while earlier stages have
outstanding receives.  Exercise the dispatcher directly so independent
EP groups initialize concurrently before the PP schedule posts P2P work.

The synthetic tensors are local to the dispatcher and no model parameter
participates in autograd.  Manager fields used by a real dispatch are
restored even if initialization fails.  Capacity mode
(`hybridep_capacity_factor`) is suspended for the synthetic dispatch: its
capacity is calibrated from the first *real* batch, whose routing skew a
balanced synthetic batch cannot stand in for, so the factor and any
calibrated capacity are restored untouched.

```python
nemo_automodel.components.moe.megatron.token_dispatcher._HybridEPManager.setup_metadata(
    routing_map: torch.Tensor,
    probs: torch.Tensor
)
```

Process routing map and probabilities to prepare dispatch metadata.

```python
nemo_automodel.components.moe.megatron.token_dispatcher._HybridEPManager.setup_metadata_from_indices(
    token_indices: torch.Tensor,
    token_probs: torch.Tensor
)
```

Convert from topk indices format to multihot routing_map format.

```python
class nemo_automodel.components.moe.megatron.token_dispatcher._HybridEPMetadataProcessor(
    num_experts: int,
    permute_fusion: bool,
    compact_routing: bool = False
)
```

**Bases:** `Module`

Graphable conversion of top-k metadata for dense or compact HybridEP dispatch.

```python
nemo_automodel.components.moe.megatron.token_dispatcher._HybridEPMetadataProcessor.forward(
    token_indices: torch.Tensor,
    token_probs: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]
```

Convert top-k metadata while preserving the graph capture boundary.

**Parameters:**

**`token_indices`** `torch.Tensor`

Expert indices with shape \[tokens, top_k]. Invalid slots contain -1.

---

**`token_probs`** `torch.Tensor`

Routing probabilities with shape \[tokens, top_k].

---

**Returns:** `torch.Tensor`

Routing metadata and dense probabilities. In compact mode, the metadata aliases

```python
nemo_automodel.components.moe.megatron.token_dispatcher._align_up(
    count: int,
    multiple: int
) -> int
```

Round `count` up to the next multiple of `multiple`.

```python
nemo_automodel.components.moe.megatron.token_dispatcher._assert_no_hybridep_overflow(
    handle,
    capacity: int
) -> None
```

Device-side guard: HybridEP truncates silently when the capacity is exceeded and sets
`overflow_flag`, the last handle item; fail loudly instead of training on dropped tokens.

**Parameters:**

**`handle`**

HybridEP dispatch handle (tuple). Its last item is the one-element device
`overflow_flag` tensor, non-zero after a truncated dispatch. DeepEP keeps the flag last
precisely so callers read it as `handle[-1]` (the slot count changed at 10d4dd7, which
inserted `num_of_valid_tokens` before it).

---

**`capacity`** `int`

Permuted-row capacity the dispatch buffers were sized to (quoted in the message).

---

```python
nemo_automodel.components.moe.megatron.token_dispatcher._HYBRIDEP_TOKEN_ALIGNMENT = 4
```

```python
nemo_automodel.components.moe.megatron.token_dispatcher._STATIC_ROUTING_PAD_PIN = os.environ.get('NEMO_STATIC_ROUTING_PAD_PIN', '1') != '0'
```