core.transformer.moe.megakernel.backend

Common interface implemented by MoE megakernel backends.

Module Contents

Classes

MegakernelBackend

Backend-neutral interface for replacing an MCore MoE layer.

API

class core.transformer.moe.megakernel.backend.MegakernelBackend

Bases: torch.nn.Module, abc.ABC

Backend-neutral interface for replacing an MCore MoE layer.

abstractmethod forward(

hidden_states: torch.Tensor,

probs: torch.Tensor,

routing_map: torch.Tensor,

) → torch.Tensor

Execute the MoE workload using MCore’s authoritative routing result.