Struct MoeLayer 

pub struct MoeLayer<R: Runtime> { /* private fields */ }

Mixture of Experts layer.

Routes tokens to top-k experts, computes expert outputs, and returns the weighted combination.

All computation stays on-device — no GPU-CPU transfers.

Implementations

impl<R: Runtime> MoeLayer<R>

pub fn new( router: MoeRouter<R>, experts: Vec<Expert<R>>, shared_expert: Option<Expert<R>>, ) -> Self

pub fn router(&self) -> &MoeRouter<R>

pub fn experts(&self) -> &[Expert<R>]

pub fn shared_expert(&self) -> Option<&Expert<R>>

pub fn forward<C>(&self, client: &C, x: &Var<R>) -> Result<MoeOutput<R>>

where R: Runtime<DType = DType>, C: RuntimeClient<R> + TensorOps<R> + ScalarOps<R> + ReduceOps<R> + ShapeOps<R> + ActivationOps<R> + SortingOps<R> + IndexingOps<R> + CompareOps<R>, R::Client: RuntimeClient<R> + TensorOps<R> + ScalarOps<R> + ActivationOps<R> + ReduceOps<R> + ShapeOps<R>,

Forward pass with auxiliary loss.

Input: [num_tokens, hidden_size] Returns: MoeOutput with output tensor and aux_loss

Strategy: iterate over experts (not tokens). For each expert, compute output for all tokens, then mask-and-weight by routing decisions. All ops stay on-device.