nemo_automodel.components.moe.megatron.token_dispatcher

Module Contents

Classes

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

Functions

Name	Description
`_align_up`	Round `count` up to the next multiple of `multiple`.
`_assert_no_hybridep_overflow`	Device-side guard: HybridEP truncates silently when the capacity is exceeded and sets

Data

_HYBRIDEP_TOKEN_ALIGNMENT

_STATIC_ROUTING_PAD_PIN

API

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.

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.

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.

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.

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

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.

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.

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.

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.

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.

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.

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:

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

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.

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:

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

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:

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

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:

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

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.

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

A manager class to handle fused all-to-all communication processes for MoE models using DeepEP backend. See 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

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.

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

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

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

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

Get the number of tokens per expert.

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

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

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

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.

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

abstract

Combine the hidden_states after expert processing.

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.

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

abstract

Get the metadata of the dispatched hidden_states.

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.

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.

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.

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

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

num_unpadded_tokens

int | None = None

routing_map

Tensor | None = None

token_indices

Tensor | None = None

token_probs

Tensor | None = None

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.

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.

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.

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

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

boolDefaults to True

Ignored because HybridEP does not expose asynchronous completion here.

allocate_on_comm_stream

boolDefaults to True

Ignored because HybridEP owns its communication stream.

capacity_mode

bool | NoneDefaults to 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].

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

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

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

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

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.

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.

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.

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.

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

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

Round count up to the next multiple of multiple.

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

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

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