`core.fusions.fused_moe_residual_add`#

Fused MoE output residual add: residual + (expert_output + shared_expert_output).

Replaces the two elementwise adds that follow the MoE token combine (the shared-expert add in MoELayer.postprocess and the residual add of the bias-dropout-add without bias or dropout) with one kernel that reads the three tensors once. The routed and shared expert outputs are summed and rounded to the activation dtype before the residual is added, so the result is bit-identical to the two unfused adds.

Module Contents#

Classes#

_FusedMoEResidualAdd

Functions#

`_moe_residual_add_kernel`
`can_use_fused_moe_residual_add`	Return whether `fused_moe_residual_add` supports the given tensors.
`fused_moe_residual_add`	Compute `residual + (expert_output + shared_expert_output)` in one kernel.

Data#

`_BLOCK_SIZE`
`_SUPPORTED_DTYPES`

API#

core.fusions.fused_moe_residual_add._BLOCK_SIZE#: 4096

core.fusions.fused_moe_residual_add._SUPPORTED_DTYPES#: ()

core.fusions.fused_moe_residual_add._moe_residual_add_kernel( expert_output_ptr, shared_expert_output_ptr, residual_ptr, output_ptr, num_elements, BLOCK_SIZE: triton.language.constexpr, )#

class core.fusions.fused_moe_residual_add._FusedMoEResidualAdd#

Bases: torch.autograd.Function

static forward(ctx, expert_output, shared_expert_output, residual)#: Return residual + (expert_output + shared_expert_output).

static backward(ctx, grad_output)#: Pass the output gradient to all three inputs, which enter the sum with unit weight.

core.fusions.fused_moe_residual_add.can_use_fused_moe_residual_add( expert_output: torch.Tensor, shared_expert_output: torch.Tensor, residual: torch.Tensor, ) → bool#

Return whether fused_moe_residual_add supports the given tensors.

The tensors must be contiguous CUDA tensors on the same device with the same shape and the same BF16, FP16 or FP32 dtype, and Triton must be available.

core.fusions.fused_moe_residual_add.fused_moe_residual_add( expert_output: torch.Tensor, shared_expert_output: torch.Tensor, residual: torch.Tensor, ) → torch.Tensor#

Compute residual + (expert_output + shared_expert_output) in one kernel.

The result is bit-identical to the two unfused adds, and the output gradient is passed to all three inputs. The inputs must satisfy can_use_fused_moe_residual_add.

Parameters:

expert_output (torch.Tensor) – Combined routed-expert output.
shared_expert_output (torch.Tensor) – Shared-expert output.
residual (torch.Tensor) – Residual stream input of the MLP.

Returns:

The MLP output added to the residual stream.

Return type:

torch.Tensor

core.fusions.fused_moe_residual_add#