Router.ixx File Reference

Mixture-of-experts router: which experts each token runs, and with what combine weights. More...

#include <cmath>
#include <cstdint>
#include <format>
#include <memory>
#include <optional>
#include <sstream>
#include <stdexcept>
#include <string>
#include <vector>
import Serialization.SafeTensors;
import Serialization.Tensor;
import Serialization.ModelArchive;
import Compute.CpuMemoryResource;
import Compute.ExecutionContextFactory;
import Compute.ExecutionContext;
import Compute.DeviceTypeTraits;
import Compute.Observation;
import Compute.DeviceType;
import Compute.OperationTraits;
import Dnn.Components.RouterConfig;
import Dnn.Component;
import Dnn.ComponentType;
import Dnn.TensorDataType;
import Dnn.CompositeComponent;
import Compute.IExecutionContext;
import Dnn.TensorDataTypeTraits;
import Dnn.TensorOps;
import Dnn.ITensor;
import Dnn.TensorTypes;
import Dnn.Tensor;
import Dnn.Components.Linear;
import Serialization.Mode;
import Compute.DeviceId;
import Dnn.Components.RmsNorm;
import Compute.DeviceAllocation;

Classes
class	Mila::Dnn::Router< TDeviceType, TPrecision >
	Routes each token to its top_k experts. More...
struct	Mila::Dnn::Router< TDeviceType, TPrecision >::Routing
	One call's result. Both views are router-owned and valid until the next forward(). More...

Namespaces
namespace	Mila
	Mila main API namespace.

Detailed Description

Mixture-of-experts router: which experts each token runs, and with what combine weights.

Gemma 4's routing chain (Specifications/Gemma4MoE.md Phase 1): logits = proj( rms_norm_without_scale( x ) * scale * hidden_size^-0.5 ) weights, indices = top_k( softmax( logits ) ), renormalized, times per_expert_scale[ index ]