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

## Module Contents

### Classes

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

### Functions

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

### Data

[`HAVE_DEEP_EP`](#nemo_automodel-components-moe-megatron-fused_a2a-HAVE_DEEP_EP)

[`HAVE_HYBRIDEP`](#nemo_automodel-components-moe-megatron-fused_a2a-HAVE_HYBRIDEP)

[`HAVE_UCCL_EP`](#nemo_automodel-components-moe-megatron-fused_a2a-HAVE_UCCL_EP)

[`_JIT_CACHE_ROOT`](#nemo_automodel-components-moe-megatron-fused_a2a-_JIT_CACHE_ROOT)

[`_buffer`](#nemo_automodel-components-moe-megatron-fused_a2a-_buffer)

[`_dispatch_replay_state`](#nemo_automodel-components-moe-megatron-fused_a2a-_dispatch_replay_state)

[`_hybrid_ep_buffer`](#nemo_automodel-components-moe-megatron-fused_a2a-_hybrid_ep_buffer)

[`_hybrid_ep_initialized_capacity`](#nemo_automodel-components-moe-megatron-fused_a2a-_hybrid_ep_initialized_capacity)

[`_hybrid_ep_runtime_signature`](#nemo_automodel-components-moe-megatron-fused_a2a-_hybrid_ep_runtime_signature)

[`_hybridep_dispatch_replay_state`](#nemo_automodel-components-moe-megatron-fused_a2a-_hybridep_dispatch_replay_state)

[`_jit_proc_dir`](#nemo_automodel-components-moe-megatron-fused_a2a-_jit_proc_dir)

[`_jit_shared_dir`](#nemo_automodel-components-moe-megatron-fused_a2a-_jit_shared_dir)

[`_jit_stored`](#nemo_automodel-components-moe-megatron-fused_a2a-_jit_stored)

[`_nvshmem_available`](#nemo_automodel-components-moe-megatron-fused_a2a-_nvshmem_available)

[`_uccl_buffer`](#nemo_automodel-components-moe-megatron-fused_a2a-_uccl_buffer)

[`logger`](#nemo_automodel-components-moe-megatron-fused_a2a-logger)

### API

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

---

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

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

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

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

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

**Bases:** `Function`

Fused combine operation for MoE output combining computation and communication.

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

staticmethod

Backward pass of fused combine.

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

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

**Bases:** `Function`

Fused dispatch operation for MoE routing combining computation and communication.

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

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

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

**Bases:** `Function`

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

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

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

---

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

Expert outputs with shape \[dispatched\_tokens, hidden].

---

**`handle`** `object`

Opaque handle returned by dispatch.

---

**`num_permuted_tokens`** `int | None` — default: None

Optional static dispatched-token capacity.

---

**`pad_multiple`** `int | None` — default: None

Optional token padding multiple.

---

**`fuse_permute`** `bool` — default: False

Whether to fuse output unpermutation into combine.

---

**Returns:** `torch.Tensor`

Combined hidden states with shape \[tokens, hidden].

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

**Bases:** `Function`

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

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

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

---

**`x`** `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`** `int` — default: 24

Number of SMs used by the dispatch API.

---

**`num_sms_combine_api`** `int` — default: 24

Number of SMs used by the combine API.

---

**`num_permuted_tokens`** `int | None` — default: None

Optional static dispatched-token capacity.

---

**`pad_multiple`** `int | None` — default: None

Optional token padding multiple.

---

**`topk_idx`** `torch.Tensor | None` — default: None

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

---

**`num_experts`** `int | None` — default: None

Global expert count required with compact indices.

---

**`fuse_permute`** `bool` — default: False

Whether to fuse token permutation into dispatch.

---

**`num_sms_preprocessing_api`** `int | None` — default: None

Optional SM count for routing-metadata preprocessing.

---

**`num_blocks_permute`** `int | None` — default: None

Optional dispatch permutation block count.

---

**`num_blocks_unpermute`** `int | None` — default: None

Optional combine unpermutation block count.

---

**Returns:** `torch.Tensor`

Dispatched hidden states with shape \[dispatched\_tokens, hidden], grouped by local

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

---

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

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

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

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

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

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

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

**Bases:** `Function`

Fused combine using UCCL-EP instead of DeepEP.

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

staticmethod

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

staticmethod

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

**Bases:** `Function`

Fused dispatch using UCCL-EP instead of DeepEP.

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

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

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

---

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

**Bases:** `local`

**`mode`** `str | None = None`

---

**`recorder`** `HybridEPDispatchReplayRecorder | None = None`

---

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

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

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

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

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

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

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

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

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

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

**`x`**

Input tensor

---

**`group`**

Process group

---

**`handle`**

Communication handle

---

**`previous_event`**

Previous CUDA event

---

**Returns:**

Result of FusedCombine

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

**`x`**

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

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

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

Calculate the number of hidden bytes for a tensor.

**Parameters:**

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

Input tensor

---

**Returns:** `int`

Number of hidden bytes

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

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

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

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

Permuted expert outputs with shape \[dispatched\_tokens, hidden].

---

**`handle`** `object`

Opaque handle returned by `hybrid_ep_dispatch`.

---

**`num_permuted_tokens`** `int | None` — default: None

Static token capacity used by non-blocking dispatch.

---

**`pad_multiple`** `int | None` — default: None

Optional token padding multiple.

---

**`fuse_permute`** `bool` — default: False

Fuse output unpermutation into the combine kernel.

---

**Returns:** `torch.Tensor`

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

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

**`x`** `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`** `int` — default: 24

SMs reserved for HybridEP dispatch.

---

**`num_sms_combine_api`** `int` — default: 24

SMs reserved for HybridEP combine.

---

**`num_permuted_tokens`** `int | None` — default: None

Static output capacity for non-blocking dispatch.

---

**`pad_multiple`** `int | None` — default: None

Optional token padding multiple.

---

**`topk_idx`** `torch.Tensor | None` — default: None

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

---

**`num_experts`** `int | None` — default: None

Global expert count required with compact indices.

---

**`fuse_permute`** `bool` — default: False

Fuse token permutation into the dispatch kernel.

---

**`num_sms_preprocessing_api`** `int | None` — default: None

Optional SM count for routing-metadata preprocessing.

---

**`num_blocks_permute`** `int | None` — default: None

Optional dispatch permutation block count.

---

**`num_blocks_unpermute`** `int | None` — default: None

Optional combine unpermutation block count.

---

**Returns:** `torch.Tensor`

Dispatched hidden states with shape \[dispatched\_tokens, hidden], grouped by local

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

```python
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 | None` — default: None

Optional SM count for routing-metadata preprocessing.

---

**`num_blocks_permute`** `int | None` — default: None

Optional dispatch permutation block count.

---

**`num_blocks_unpermute`** `int | None` — default: None

Optional combine unpermutation block count.

---

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

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

Reset the HybridEP buffer and its published pipeline runtime state.

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

Sets the number of SMs to use for DeepEP.

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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