nemo_automodel.components.moe.megatron.fused_a2a

Module Contents

Classes

Name	Description
`DispatchReplayRecorder`	Per-checkpoint-frame log of DeepEP dispatch results, replayed on recompute.
`FusedCombine`	Fused combine operation for MoE output combining computation and communication.
`FusedDispatch`	Fused dispatch operation for MoE routing combining computation and communication.
`HybridEPCombine`	Fused combine operation for permute + combine a2a + permute using the HybridEP backend.
`HybridEPDispatch`	Fused dispatch operation for permute + dispatch a2a + permute using the HybridEP backend.
`HybridEPDispatchReplayRecorder`	Record HybridEP layouts from checkpoint forward for deterministic replay.
`UCCLFusedCombine`	Fused combine using UCCL-EP instead of DeepEP.
`UCCLFusedDispatch`	Fused dispatch using UCCL-EP instead of DeepEP.
`_DispatchReplayState`	Thread-local replay binding. `threading.local` runs `__init__` once per
`_HybridEPDispatchReplayState`	-

Functions

Name	Description
`_active_recorder`	-
`_deep_ep_version`	Installed DeepEP version (dist-info first: the package itself carries no `__version__`).
`_hybrid_ep_jit_dirs`	Return (shared cache dir for this toolchain and GPU, this process’s HybridEP JIT dir).
`_is_nvshmem_available`	Check if DeepEP was compiled with NVSHMEM support.
`_replay_mode`	-
`_warm_start_hybrid_ep_jit`	Copy the shared cache’s kernels into this process’s JIT dir. Returns whether the cache is active.
`dispatch_replay_scope`	Bind `recorder` for the enclosed region. `mode` is ‘record’ or ‘replay’.
`free_buffer`	Destroy the global DeepEP `Buffer` and release its NVSHMEM/cpp runtime.
`fused_combine`	Perform fused combine operation if deep_ep is available.
`fused_dispatch`	Perform fused dispatch operation if deep_ep is available.
`get_buffer`	Get or create a buffer for all-to-all communication.
`get_hidden_bytes`	Calculate the number of hidden bytes for a tensor.
`get_hybrid_ep_initialized_capacity`	Return process-global initialized capacity after validating the resource signature.
`get_uccl_buffer`	Get or create a UCCL-EP buffer for all-to-all communication.
`hybrid_ep_combine`	Combine expert outputs through HybridEP.
`hybrid_ep_dispatch`	Dispatch tokens through HybridEP while retaining dense probability layout.
`hybridep_dispatch_replay_scope`	Bind a HybridEP dispatch recorder for checkpoint forward or recompute.
`init_hybrid_ep_buffer`	Initialize the HybridEP buffer, including buffer allocation and metadata initialization.
`publish_hybrid_ep_initialized_capacity`	Publish a successfully initialized process-global HybridEP capacity.
`reset_hybrid_ep_buffer`	Reset the HybridEP buffer and its published pipeline runtime state.
`set_deepep_num_sms`	Sets the number of SMs to use for DeepEP.
`set_uccl_num_sms`	Sets the number of SMs to use for UCCL-EP.
`store_hybrid_ep_jit_cache`	Store the kernels HybridEP compiled in this process into the shared cache (local rank 0 only).
`uccl_fused_combine`	Perform fused combine using UCCL-EP.
`uccl_fused_dispatch`	Perform fused dispatch using UCCL-EP.

Data

_dispatch_replay_state

_hybrid_ep_buffer

_hybrid_ep_initialized_capacity

_hybrid_ep_runtime_signature

_hybridep_dispatch_replay_state

API

class nemo_automodel.components.moe.megatron.fused_a2a.DispatchReplayRecorder()

Per-checkpoint-frame log of DeepEP dispatch results, replayed on recompute.

_cursor

= 0

_records

list = []

replay_misses

= 0

nemo_automodel.components.moe.megatron.fused_a2a.DispatchReplayRecorder.record(
    entry
) -> None

nemo_automodel.components.moe.megatron.fused_a2a.DispatchReplayRecorder.rewind() -> None

nemo_automodel.components.moe.megatron.fused_a2a.DispatchReplayRecorder.take()

Next recorded dispatch, or None when the replay outruns the log.

class nemo_automodel.components.moe.megatron.fused_a2a.FusedCombine()

Bases: Function

Fused combine operation for MoE output combining computation and communication.

nemo_automodel.components.moe.megatron.fused_a2a.FusedCombine.backward(
    ctx,
    grad_output,
    previous_event = None
)

staticmethod

Backward pass of fused combine.

nemo_automodel.components.moe.megatron.fused_a2a.FusedCombine.forward(
    ctx,
    x,
    group,
    handle,
    async_finish = False,
    allocate_on_comm_stream = False
)

staticmethod

Forward pass of fused combine.

class nemo_automodel.components.moe.megatron.fused_a2a.FusedDispatch()

Bases: Function

Fused dispatch operation for MoE routing combining computation and communication.

nemo_automodel.components.moe.megatron.fused_a2a.FusedDispatch.backward(
    ctx,
    grad_output,
    grad_token_indices,
    grad_token_probs,
    grad_tokens_per_expert,
    grad_handle
)

staticmethod

Backward pass of fused dispatch.

nemo_automodel.components.moe.megatron.fused_a2a.FusedDispatch.forward(
    ctx,
    x,
    token_indices,
    token_probs,
    num_experts,
    group,
    async_finish = False,
    allocate_on_comm_stream = False
)

staticmethod

Forward pass of fused dispatch.

class nemo_automodel.components.moe.megatron.fused_a2a.HybridEPCombine()

Bases: Function

Fused combine operation for permute + combine a2a + permute using the HybridEP backend.

nemo_automodel.components.moe.megatron.fused_a2a.HybridEPCombine.backward(
    ctx: torch.autograd.function.FunctionCtx,
    grad_x: torch.Tensor
) -> tuple[torch.Tensor, None, None, None, None]

staticmethod

Redispatch combined-output gradients through the saved layout.

Parameters:

ctx

torch.autograd.function.FunctionCtx

Autograd context populated by forward.

grad_x

torch.Tensor

Combined-output gradients with shape [tokens, hidden].

Returns: torch.Tensor

Gradient for x with shape [dispatched_tokens, hidden], followed by None

nemo_automodel.components.moe.megatron.fused_a2a.HybridEPCombine.forward(
    ctx: torch.autograd.function.FunctionCtx,
    x: torch.Tensor,
    handle: object,
    num_permuted_tokens: int | None = None,
    pad_multiple: int | None = None,
    fuse_permute: bool = False
) -> torch.Tensor

staticmethod

Combine dispatched expert outputs through HybridEP.

Parameters:

ctx

torch.autograd.function.FunctionCtx

Autograd context that retains the dispatch handle for backward.

torch.Tensor

Expert outputs with shape [dispatched_tokens, hidden].

handle

object

Opaque handle returned by dispatch.

num_permuted_tokens

int | NoneDefaults to None

Optional static dispatched-token capacity.

pad_multiple

int | NoneDefaults to None

Optional token padding multiple.

fuse_permute

boolDefaults to False

Whether to fuse output unpermutation into combine.

Returns: torch.Tensor

Combined hidden states with shape [tokens, hidden].

class nemo_automodel.components.moe.megatron.fused_a2a.HybridEPDispatch()

Bases: Function

Fused dispatch operation for permute + dispatch a2a + permute using the HybridEP backend.

nemo_automodel.components.moe.megatron.fused_a2a.HybridEPDispatch.backward(
    ctx: torch.autograd.function.FunctionCtx,
    grad_x: torch.Tensor,
    grad_probs: torch.Tensor,
    grad_scaling_factor: torch.Tensor | None,
    grad_tokens_per_expert: torch.Tensor | None,
    grad_handle: None
) -> tuple[torch.Tensor | None, ...]

staticmethod

Combine gradients while retaining the dense probability-gradient layout.

Parameters:

ctx

torch.autograd.function.FunctionCtx

Autograd context populated by forward.

grad_x

torch.Tensor

Hidden-state gradients with shape [dispatched_tokens, hidden].

grad_probs

torch.Tensor

Probability gradients with shape [dispatched_tokens], aligned with grad_x.

grad_scaling_factor

torch.Tensor | None

Optional gradient for dispatch scaling metadata.

grad_tokens_per_expert

torch.Tensor | None

Ignored gradient for token counts with shape [local_experts].

grad_handle

None

Ignored gradient for the opaque dispatch handle.

Returns: torch.Tensor | None

Gradients for every forward argument. The hidden-state gradient has shape

nemo_automodel.components.moe.megatron.fused_a2a.HybridEPDispatch.forward(
    ctx: torch.autograd.function.FunctionCtx,
    x: torch.Tensor,
    routing_map: torch.Tensor | None,
    probs: torch.Tensor,
    group: torch.distributed.ProcessGroup,
    num_local_experts: int,
    num_sms_dispatch_api: int = 24,
    num_sms_combine_api: int = 24,
    num_permuted_tokens: int | None = None,
    pad_multiple: int | None = None,
    topk_idx: torch.Tensor | None = None,
    num_experts: int | None = None,
    fuse_permute: bool = False,
    num_sms_preprocessing_api: int | None = None,
    num_blocks_permute: int | None = None,
    num_blocks_unpermute: int | None = None
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor, object]

staticmethod

Dispatch hidden states using compact indices or a dense routing map.

The probability tensor remains dense so its gradient layout is unchanged. When both routing representations are provided, HybridEP gives routing_map precedence.

Parameters:

ctx

torch.autograd.function.FunctionCtx

Autograd context that retains the dispatch handle for backward.

torch.Tensor

Hidden states with shape [tokens, hidden].

routing_map

torch.Tensor | None

Optional Boolean routing map with shape [tokens, experts].

probs

torch.Tensor

Dense routing probabilities with shape [tokens, experts].

group

torch.distributed.ProcessGroup

Expert-parallel process group.

num_local_experts

int

Number of experts owned by each rank.

num_sms_dispatch_api

intDefaults to 24

Number of SMs used by the dispatch API.

num_sms_combine_api

intDefaults to 24

Number of SMs used by the combine API.

num_permuted_tokens

int | NoneDefaults to None

Optional static dispatched-token capacity.

pad_multiple

int | NoneDefaults to None

Optional token padding multiple.

topk_idx

torch.Tensor | NoneDefaults to None

Optional global expert indices with shape [tokens, top_k], with -1 for unused slots.

num_experts

int | NoneDefaults to None

Global expert count required with compact indices.

fuse_permute

boolDefaults to False

Whether to fuse token permutation into dispatch.

num_sms_preprocessing_api

int | NoneDefaults to None

Optional SM count for routing-metadata preprocessing.

num_blocks_permute

int | NoneDefaults to None

Optional dispatch permutation block count.

num_blocks_unpermute

int | NoneDefaults to None

Optional combine unpermutation block count.

Returns: torch.Tensor

Dispatched hidden states with shape [dispatched_tokens, hidden], grouped by local

class nemo_automodel.components.moe.megatron.fused_a2a.HybridEPDispatchReplayRecorder()

Record HybridEP layouts from checkpoint forward for deterministic replay.

HybridEP’s sync-free replay API expects the dispatched-token extent as a host integer. Reading it with .item() right after every checkpointed block stalls the host until the GPU has drained that block, once per layer per microbatch; on a pipeline-parallel step those stalls surface as idle GPU time. The recorder therefore splits the read in two: finalize (called when the checkpoint-forward op context exits) runs the reduction and starts an asynchronous device-to-host copy, and resolve (called when the recompute context opens, long after the copy completed) turns the pinned scalars into integers. Neither step runs inside a checkpoint op context, so the forward and recompute traces stay identical.

_copies

list = []

_cursor

= 0

_records

list = []

replay_misses

= 0

nemo_automodel.components.moe.megatron.fused_a2a.HybridEPDispatchReplayRecorder.finalize() -> None

Start the device-to-host copy of each layout extent (no host synchronization).

Runs after the checkpoint-forward op context exits, so the reduction does not enter the forward trace; CPU tensors are converted immediately.

nemo_automodel.components.moe.megatron.fused_a2a.HybridEPDispatchReplayRecorder.record(
    handle,
    tokens_per_expert
) -> None

nemo_automodel.components.moe.megatron.fused_a2a.HybridEPDispatchReplayRecorder.resolve() -> None

Turn the copied extents into host integers, waiting only for copies still in flight.

Called when the recompute context opens — outside its op context, so the conversion never enters the recompute trace.

nemo_automodel.components.moe.megatron.fused_a2a.HybridEPDispatchReplayRecorder.rewind() -> None

nemo_automodel.components.moe.megatron.fused_a2a.HybridEPDispatchReplayRecorder.take()

Return the next forward dispatch record, or None on divergence.

class nemo_automodel.components.moe.megatron.fused_a2a.UCCLFusedCombine()

Bases: Function

Fused combine using UCCL-EP instead of DeepEP.

nemo_automodel.components.moe.megatron.fused_a2a.UCCLFusedCombine.backward(
    ctx,
    grad_output,
    _grad_event = None
)

staticmethod

nemo_automodel.components.moe.megatron.fused_a2a.UCCLFusedCombine.forward(
    ctx,
    x,
    group,
    handle,
    async_finish = False,
    allocate_on_comm_stream = False
)

staticmethod

class nemo_automodel.components.moe.megatron.fused_a2a.UCCLFusedDispatch()

Bases: Function

Fused dispatch using UCCL-EP instead of DeepEP.

nemo_automodel.components.moe.megatron.fused_a2a.UCCLFusedDispatch.backward(
    ctx,
    grad_output,
    grad_token_indices,
    grad_token_probs,
    grad_tokens_per_expert,
    grad_handle
)

staticmethod

nemo_automodel.components.moe.megatron.fused_a2a.UCCLFusedDispatch.forward(
    ctx,
    x,
    token_indices,
    token_probs,
    num_experts,
    group,
    async_finish = False,
    allocate_on_comm_stream = False
)

staticmethod

class nemo_automodel.components.moe.megatron.fused_a2a._DispatchReplayState()

Bases: local

Thread-local replay binding. threading.local runs __init__ once per thread, so both attributes always exist on every thread that touches it.

mode

str | None = None

recorder

DispatchReplayRecorder | None = None

class nemo_automodel.components.moe.megatron.fused_a2a._HybridEPDispatchReplayState()

Bases: local

mode

str | None = None

recorder

HybridEPDispatchReplayRecorder | None = None

nemo_automodel.components.moe.megatron.fused_a2a._active_recorder() -> nemo_automodel.components.moe.megatron.fused_a2a.DispatchReplayRecorder | None

nemo_automodel.components.moe.megatron.fused_a2a._deep_ep_version() -> str

Installed DeepEP version (dist-info first: the package itself carries no __version__).

nemo_automodel.components.moe.megatron.fused_a2a._hybrid_ep_jit_dirs() -> tuple[str, str]

Return (shared cache dir for this toolchain and GPU, this process’s HybridEP JIT dir).

The per-process path mirrors get_jit_dir() in DeepEP’s csrc/hybrid_ep/jit/compiler.cu.

nemo_automodel.components.moe.megatron.fused_a2a._is_nvshmem_available() -> bool

Check if DeepEP was compiled with NVSHMEM support.

Uses is_sm90_compiled() as proxy — DeepEP’s build enforces that NVSHMEM is disabled when SM90 features are disabled.

nemo_automodel.components.moe.megatron.fused_a2a._replay_mode() -> str | None

nemo_automodel.components.moe.megatron.fused_a2a._warm_start_hybrid_ep_jit() -> bool

Copy the shared cache’s kernels into this process’s JIT dir. Returns whether the cache is active.

nemo_automodel.components.moe.megatron.fused_a2a.dispatch_replay_scope(
    recorder: nemo_automodel.components.moe.megatron.fused_a2a.DispatchReplayRecorder | None,
    mode: str
)

Bind recorder for the enclosed region. mode is ‘record’ or ‘replay’.

nemo_automodel.components.moe.megatron.fused_a2a.free_buffer() -> None

Destroy the global DeepEP Buffer and release its NVSHMEM/cpp runtime.

DeepEP keeps a process-global communication buffer backed by NVSHMEM symmetric memory. It is normally never torn down (destroy_process_group hangs on DeepEP’s NCCL sub-groups, so cleanup is skipped), but that leftover GPU state survives process exit for the whole Slurm allocation and corrupts subsequent forwards. Destroying the buffer first frees the runtime and lets a clean destroy_process_group follow without hanging.

nemo_automodel.components.moe.megatron.fused_a2a.fused_combine(
    x,
    group,
    handle,
    async_finish = False,
    allocate_on_comm_stream = False
)

Perform fused combine operation if deep_ep is available.

Parameters:

Input tensor

group

Process group

handle

Communication handle

previous_event

Previous CUDA event

Returns:

Result of FusedCombine

nemo_automodel.components.moe.megatron.fused_a2a.fused_dispatch(
    x,
    token_indices,
    token_probs,
    num_experts,
    group,
    async_finish = False,
    allocate_on_comm_stream = False
)

Perform fused dispatch operation if deep_ep is available.

Parameters:

Input tensor [num_tokens, hidden_size]

token_indices

Token routing indices [num_tokens, topk]

token_probs

Token routing probabilities [num_tokens, topk]

num_experts

Number of experts

group

Process group

previous_event

Previous CUDA event

Returns:

Result of FusedDispatch

nemo_automodel.components.moe.megatron.fused_a2a.get_buffer(
    group: torch.distributed.ProcessGroup,
    hidden_bytes: int
)

Get or create a buffer for all-to-all communication.

Parameters:

group

torch.distributed.ProcessGroup

Process group for communication

hidden_bytes

int

Number of hidden bytes needed

Returns:

Communication buffer

nemo_automodel.components.moe.megatron.fused_a2a.get_hidden_bytes(
    x: torch.Tensor
) -> int

Calculate the number of hidden bytes for a tensor.

Parameters:

torch.Tensor

Input tensor

Returns: int

Number of hidden bytes

nemo_automodel.components.moe.megatron.fused_a2a.get_hybrid_ep_initialized_capacity(
    signature: tuple[object, ...]
) -> int

Return process-global initialized capacity after validating the resource signature.

Parameters:

signature

tuple[object, ...]

Full configuration and device signature for the HybridEP buffer.

Returns: int

The greatest successfully initialized token capacity, or zero before initialization.

Raises:

RuntimeError: If the active process-global buffer has an incompatible signature.

nemo_automodel.components.moe.megatron.fused_a2a.get_uccl_buffer(
    group: torch.distributed.ProcessGroup,
    hidden_bytes: int
)

Get or create a UCCL-EP buffer for all-to-all communication.

nemo_automodel.components.moe.megatron.fused_a2a.hybrid_ep_combine(
    x: torch.Tensor,
    handle: object,
    num_permuted_tokens: int | None = None,
    pad_multiple: int | None = None,
    fuse_permute: bool = False
) -> torch.Tensor

Combine expert outputs through HybridEP.

Parameters:

torch.Tensor

Permuted expert outputs with shape [dispatched_tokens, hidden].

handle

object

Opaque handle returned by hybrid_ep_dispatch.

num_permuted_tokens

int | NoneDefaults to None

Static token capacity used by non-blocking dispatch.

pad_multiple

int | NoneDefaults to None

Optional token padding multiple.

fuse_permute

boolDefaults to False

Fuse output unpermutation into the combine kernel.

Returns: torch.Tensor

Combined hidden states in the original token order with shape [tokens, hidden].

nemo_automodel.components.moe.megatron.fused_a2a.hybrid_ep_dispatch(
    x: torch.Tensor,
    routing_map: torch.Tensor | None,
    probs: torch.Tensor,
    group: torch.distributed.ProcessGroup,
    num_local_experts: int,
    num_sms_dispatch_api: int = 24,
    num_sms_combine_api: int = 24,
    num_permuted_tokens: int | None = None,
    pad_multiple: int | None = None,
    topk_idx: torch.Tensor | None = None,
    num_experts: int | None = None,
    fuse_permute: bool = False,
    num_sms_preprocessing_api: int | None = None,
    num_blocks_permute: int | None = None,
    num_blocks_unpermute: int | None = None
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor, object]

Dispatch tokens through HybridEP while retaining dense probability layout.

Parameters:

torch.Tensor

Hidden states with shape [tokens, hidden].

routing_map

torch.Tensor | None

Optional dense routing map with shape [tokens, experts].

probs

torch.Tensor

Dense routing probabilities with shape [tokens, experts].

group

torch.distributed.ProcessGroup

Expert-parallel process group.

num_local_experts

int

Number of experts owned by each rank.

num_sms_dispatch_api

intDefaults to 24

SMs reserved for HybridEP dispatch.

num_sms_combine_api

intDefaults to 24

SMs reserved for HybridEP combine.

num_permuted_tokens

int | NoneDefaults to None

Static output capacity for non-blocking dispatch.

pad_multiple

int | NoneDefaults to None

Optional token padding multiple.

topk_idx

torch.Tensor | NoneDefaults to None

Optional global expert indices with shape [tokens, top_k], with -1 for unused slots.

num_experts

int | NoneDefaults to None

Global expert count required with compact indices.

fuse_permute

boolDefaults to False

Fuse token permutation into the dispatch kernel.

num_sms_preprocessing_api

int | NoneDefaults to None

Optional SM count for routing-metadata preprocessing.

num_blocks_permute

int | NoneDefaults to None

Optional dispatch permutation block count.

num_blocks_unpermute

int | NoneDefaults to None

Optional combine unpermutation block count.

Returns: torch.Tensor

Dispatched hidden states with shape [dispatched_tokens, hidden], grouped by local

nemo_automodel.components.moe.megatron.fused_a2a.hybridep_dispatch_replay_scope(
    recorder: nemo_automodel.components.moe.megatron.fused_a2a.HybridEPDispatchReplayRecorder | None,
    mode: str
)

Bind a HybridEP dispatch recorder for checkpoint forward or recompute.

nemo_automodel.components.moe.megatron.fused_a2a.init_hybrid_ep_buffer(
    group: torch.distributed.ProcessGroup,
    hidden_dim: int,
    seq_len: int,
    num_local_experts: int,
    num_sms_dispatch_api: int,
    num_sms_combine_api: int,
    fp8_dispatch: bool,
    num_sms_preprocessing_api: int | None = None,
    num_blocks_permute: int | None = None,
    num_blocks_unpermute: int | None = None
) -> None

Initialize the HybridEP buffer, including buffer allocation and metadata initialization.

If a runtime dispatch/combine requires a larger buffer than the one initialized, the buffer will be reallocated at runtime, incuring extra run-time overhead.

Parameters:

group

torch.distributed.ProcessGroup

Process group for HybridEP all-to-all communication.

hidden_dim

int

Hidden dimension of the input tensor.

seq_len

int

Maximum sequence length of the input tensor.

num_local_experts

int

Number of local experts.

num_sms_dispatch_api

int

Number of SMs used by the dispatch API.

num_sms_combine_api

int

Number of SMs used by the combine API.

fp8_dispatch

bool

Whether to use FP8 communication during the dispatch phase.

num_sms_preprocessing_api

int | NoneDefaults to None

Optional SM count for routing-metadata preprocessing.

num_blocks_permute

int | NoneDefaults to None

Optional dispatch permutation block count.

num_blocks_unpermute

int | NoneDefaults to None

Optional combine unpermutation block count.

nemo_automodel.components.moe.megatron.fused_a2a.publish_hybrid_ep_initialized_capacity(
    signature: tuple[object, ...],
    capacity: int
) -> None

Publish a successfully initialized process-global HybridEP capacity.

Parameters:

signature

tuple[object, ...]

Full configuration and device signature for the HybridEP buffer.

capacity

int

Successfully initialized token capacity.

Raises:

RuntimeError: If the active process-global buffer has an incompatible signature.

nemo_automodel.components.moe.megatron.fused_a2a.reset_hybrid_ep_buffer()

Reset the HybridEP buffer and its published pipeline runtime state.

nemo_automodel.components.moe.megatron.fused_a2a.set_deepep_num_sms(
    num_sms
)

Sets the number of SMs to use for DeepEP.

nemo_automodel.components.moe.megatron.fused_a2a.set_uccl_num_sms(
    num_sms
)

Sets the number of SMs to use for UCCL-EP.

nemo_automodel.components.moe.megatron.fused_a2a.store_hybrid_ep_jit_cache() -> int

Store the kernels HybridEP compiled in this process into the shared cache (local rank 0 only).

Files are written to a temporary name and atomically renamed, so concurrent writers from several nodes leave a complete file behind; existing entries are kept.

nemo_automodel.components.moe.megatron.fused_a2a.uccl_fused_combine(
    x,
    group,
    handle,
    async_finish = False,
    allocate_on_comm_stream = False
)

Perform fused combine using UCCL-EP.

nemo_automodel.components.moe.megatron.fused_a2a.uccl_fused_dispatch(
    x,
    token_indices,
    token_probs,
    num_experts,
    group,
    async_finish = False,
    allocate_on_comm_stream = False
)

Perform fused dispatch using UCCL-EP.

nemo_automodel.components.moe.megatron.fused_a2a.HAVE_DEEP_EP = True

nemo_automodel.components.moe.megatron.fused_a2a.HAVE_HYBRIDEP = True

nemo_automodel.components.moe.megatron.fused_a2a.HAVE_UCCL_EP = True

nemo_automodel.components.moe.megatron.fused_a2a._JIT_CACHE_ROOT = os.environ.get('NEMO_HYBRIDEP_JIT_CACHE')

nemo_automodel.components.moe.megatron.fused_a2a._buffer = None

nemo_automodel.components.moe.megatron.fused_a2a._dispatch_replay_state = _DispatchReplayState()

nemo_automodel.components.moe.megatron.fused_a2a._hybrid_ep_buffer = None

nemo_automodel.components.moe.megatron.fused_a2a._hybrid_ep_initialized_capacity = 0

nemo_automodel.components.moe.megatron.fused_a2a._hybrid_ep_runtime_signature: tuple[object, ...] | None = None

nemo_automodel.components.moe.megatron.fused_a2a._hybridep_dispatch_replay_state = _HybridEPDispatchReplayState()

nemo_automodel.components.moe.megatron.fused_a2a._jit_proc_dir: str | None = None

nemo_automodel.components.moe.megatron.fused_a2a._jit_shared_dir: str | None = None

nemo_automodel.components.moe.megatron.fused_a2a._jit_stored = False

nemo_automodel.components.moe.megatron.fused_a2a._nvshmem_available = None

nemo_automodel.components.moe.megatron.fused_a2a._uccl_buffer = None

nemo_automodel.components.moe.megatron.fused_a2a.logger = logging.getLogger(__name__)