core.transformer.moe.megakernel.factory

Construction entry point for MoE megakernel backends.

Module Contents

Functions

build_megakernel_backend	Build the configured megakernel backend.
prepare_megakernel_shared_expert_config	Return the config used to construct backend-compatible shared experts.
megakernel_shared_expert_init_context	Return the backend-specific context used to construct shared experts.

API

core.transformer.moe.megakernel.factory.build_megakernel_backend(

*,

config: megatron.core.transformer.transformer_config.TransformerConfig,

ep_group: torch.distributed.ProcessGroup,

routed_experts: torch.nn.Module,

shared_experts: torch.nn.Module,

num_local_experts: int,

) → megatron.core.transformer.moe.megakernel.backend.MegakernelBackend

Build the configured megakernel backend.

Backend imports are deliberately lazy so ordinary MCore MoE configurations do not acquire optional megakernel dependencies.

core.transformer.moe.megakernel.factory.prepare_megakernel_shared_expert_config(

config: megatron.core.transformer.transformer_config.TransformerConfig,

) → megatron.core.transformer.transformer_config.TransformerConfig

Return the config used to construct backend-compatible shared experts.

core.transformer.moe.megakernel.factory.megakernel_shared_expert_init_context(

config: megatron.core.transformer.transformer_config.TransformerConfig,

)

Return the backend-specific context used to construct shared experts.