`core.transformer.moe.fused_a2a`#

Module Contents#

Classes#

`FusedDispatch`	Fused dispatch operation for MoE routing combining computation and communication.
`FusedCombine`	Fused combine operation for MoE output combining computation and communication.
`HybridEPDispatch`	Fused dispatch operation for permute + dispatch a2a + permute using the HybridEP backend
`HybridEPCombine`	Fused combine operation for permute + combine a2a + permute using the HybridEP backend

Functions#

`get_hidden_bytes`	Calculate the number of hidden bytes for a tensor.
`get_buffer`	Get or create a buffer for all-to-all communication.
`init_hybrid_ep_buffer`	Initialize the HybridEP buffer, including buffer allocation and metadata initialization.
`reset_hybrid_ep_buffer`	Reset the HybridEP buffer
`ensure_nccl_ep_bootstrapped`	Initialize the process-wide NCCL EP context once. Idempotent.
`nccl_ep_finalize`	Tear down the NCCL EP context. Idempotent; safe when never bootstrapped.

Data#

`_buffer`
`_hybrid_ep_buffer`
`HYBRIDEP_TOKEN_ALIGNMENT`

API#

core.transformer.moe.fused_a2a._buffer#: None

core.transformer.moe.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:: Number of hidden bytes
Return type:: int

core.transformer.moe.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

Return type:

Buffer

class core.transformer.moe.fused_a2a.FusedDispatch#

Bases: torch.autograd.Function

Fused dispatch operation for MoE routing combining computation and communication.

static forward( ctx, x, token_indices, token_probs, num_experts, group, async_finish=False, allocate_on_comm_stream=False, )#: Forward pass of fused dispatch.

static backward( ctx, grad_output, grad_token_indices, grad_token_probs, grad_tokens_per_expert, grad_handle, )#: Backward pass of fused dispatch.

class core.transformer.moe.fused_a2a.FusedCombine#

Bases: torch.autograd.Function

Fused combine operation for MoE output combining computation and communication.

static forward( ctx, x, group, handle, async_finish=False, allocate_on_comm_stream=False, )#: Forward pass of fused combine.

static backward(ctx, grad_output, previous_event=None)#: Backward pass of fused combine.

core.transformer.moe.fused_a2a._hybrid_ep_buffer#: None

core.transformer.moe.fused_a2a.HYBRIDEP_TOKEN_ALIGNMENT#: 64

core.transformer.moe.fused_a2a.init_hybrid_ep_buffer( group: torch.distributed.ProcessGroup, hidden_dim: int, num_tokens: int, num_local_experts: int, num_sms_dispatch_api: Optional[int] = None, num_sms_combine_api: Optional[int] = None, num_blocks_permute: Optional[int] = None, num_blocks_unpermute: Optional[int] = None, fp8_dispatch: bool = False, num_sms_preprocessing_api: Optional[int] = 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.
num_tokens (int) – Maximum token count of the input tensor.
num_local_experts (int) – Number of local experts.
num_sms_dispatch_api (Optional[int]) – Number of SMs used by the dispatch API.
num_sms_combine_api (Optional[int]) – Number of SMs used by the combine API.
num_blocks_permute (Optional[int]) – Number of blocks used by the permute part.
num_blocks_unpermute (Optional[int]) – Number of blocks used by the unpermute part.
fp8_dispatch (bool) – Whether to use FP8 communication during the dispatch phase.
num_sms_preprocessing_api (Optional[int]) – Number of SMs used by the preprocessing (metadata scan) kernel.

core.transformer.moe.fused_a2a.reset_hybrid_ep_buffer()#: Reset the HybridEP buffer

class core.transformer.moe.fused_a2a.HybridEPDispatch#

Bases: torch.autograd.Function

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

static forward( ctx, x, routing_map, probs, group, num_local_experts, num_sms_dispatch_api=None, num_sms_combine_api=None, num_blocks_permute=None, num_blocks_unpermute=None, fused=False, num_permuted_tokens=None, pad_multiple=None, num_sms_preprocessing_api=108, )#: Forward pass of fused dispatch of the HybridEP backend

static backward( ctx, grad_x, grad_probs, grad_scaling_factor, grad_tokens_per_expert, grad_handle, )#: Backward pass of fused dispatch of the HybridEP backend

class core.transformer.moe.fused_a2a.HybridEPCombine#

Bases: torch.autograd.Function

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

static forward( ctx, x, handle, num_permuted_tokens=None, pad_multiple=None, fused=False, )#: Forward pass of fused combine of the HybridEP backend

static backward(ctx, grad_x)#: Backward pass of fused combine of the HybridEP backend

core.transformer.moe.fused_a2a.ensure_nccl_ep_bootstrapped( ep_group, num_experts, max_tokens_per_rank, recv_capacity_per_rank, hidden_dim, num_sms=0, zero_copy=False, )#

Initialize the process-wide NCCL EP context once. Idempotent.

Collective on ep_group: TE’s ep_bootstrap issues a barrier and borrows the group’s NCCL communicator, so every rank must call this with identical arguments before the first dispatch. Reuses TransformerEngine’s own one-time flag, so repeated calls (e.g. once per MoE layer) are no-ops.

Parameters:

ep_group (torch.distributed.ProcessGroup) – The expert-parallel process group.
num_experts (int) – Total experts across ep_group (global, not per-rank).
max_tokens_per_rank (int) – Upper bound on local input tokens per forward. Must be even (NCCL EP requires num_tokens_per_rank * inner_dim % 4 == 0).
recv_capacity_per_rank (int) – Per-rank receive-buffer capacity in tokens. Must be >= max_tokens_per_rank; runtime overflow hard-traps (no soft drop).
hidden_dim (int) – Token hidden size.
num_sms (int) – SM cap passed to TE as max_num_sms (0 lets TE/NCCL choose).

core.transformer.moe.fused_a2a.nccl_ep_finalize()#

Tear down the NCCL EP context. Idempotent; safe when never bootstrapped.

Releases the borrowed NCCL communicator and must run before the process group is destroyed.

core.transformer.moe.fused_a2a#

Module Contents#

Classes#

Functions#

Data#

API#

`core.transformer.moe.fused_a2a`#