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

### 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_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_compact_routing`** `bool = False`

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

---

**`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`**

---

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

---

**`token_probs`** `Tensor | None = 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.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.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
)
```

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

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

---

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

---

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

---

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

---

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

---

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

---

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

```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._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'
```