oa::Moe

Moe class in the OA Ml public surface.

ClC++ public headers

Inheritance

public Module

Public Methods

const oa::Matrix & oa::Moe::auxLoss()

oa::F32 oa::Moe::balanceRate()

oa::I32 oa::Moe::dFfn()

oa::I32 oa::Moe::dModel()

oa::I32 oa::Moe::expertsPerToken()

oa::Matrix oa::Moe::forward(const oa::Matrix & inX)

void oa::Moe::init(oa::I32 inDModel, oa::I32 inDFF, oa::I32 inNumExperts, oa::I32 inExpertsPerToken, oa::F32 inRmsEps = 1e-5f, oa::I32 inNumSharedExperts = 0)

const oa::Matrix & oa::Moe::lastSelectionMask()

oa::I32 oa::Moe::numExperts()

oa::I32 oa::Moe::numSharedExperts()

MoeRouteStats oa::Moe::routeStats()

oa::F32 oa::Moe::routingBias(oa::I32 inExpert)

void oa::Moe::setAuxLossAlpha(oa::F32 inAlpha)

void oa::Moe::setBalanceRate(oa::F32 inGamma)

void oa::Moe::setRouterZLossBeta(oa::F32 inBeta)

void oa::Moe::setSparseExecution(bool inEnabled)

bool oa::Moe::sparseExecution()

void oa::Moe::updateRoutingBias()

Constructor & Destructor Documentation

oa::Moe::Moe()

No public source comment is attached to this declaration.

oa::Moe::Moe(
	oa::I32 inDModel,
	oa::I32 inDFF,
	oa::I32 inNumExperts,
	oa::I32 inExpertsPerToken,
	oa::F32 inRmsEps = 1e-5f,
	oa::I32 inNumSharedExperts = 0
)

No public source comment is attached to this declaration.

Parameters

inDModel

oa::I32

—

inDFF

oa::I32

—

inNumExperts

oa::I32

—

inExpertsPerToken

oa::I32

—

inRmsEps

oa::F32

Default: 1e-5f

inNumSharedExperts

oa::I32

Default: 0

Public Method Documentation

const oa::Matrix & oa::Moe::auxLoss()

No public source comment is attached to this declaration.

Returns

const oa::Matrix &

The declared return value.

oa::F32 oa::Moe::balanceRate()

No public source comment is attached to this declaration.

Returns

oa::F32

The declared return value.

oa::I32 oa::Moe::dFfn()

No public source comment is attached to this declaration.

Returns

oa::I32

The declared return value.

oa::I32 oa::Moe::dModel()

No public source comment is attached to this declaration.

Returns

oa::I32

The declared return value.

oa::I32 oa::Moe::expertsPerToken()

No public source comment is attached to this declaration.

Returns

oa::I32

The declared return value.

oa::Matrix oa::Moe::forward(
	const oa::Matrix & inX
)

forward: x[T,D] → moe_out[T,D]

Parameters

inX

const oa::Matrix &

—

Returns

oa::Matrix

The declared return value.

void oa::Moe::init(
	oa::I32 inDModel,
	oa::I32 inDFF,
	oa::I32 inNumExperts,
	oa::I32 inExpertsPerToken,
	oa::F32 inRmsEps = 1e-5f,
	oa::I32 inNumSharedExperts = 0
)

No public source comment is attached to this declaration.

Parameters

inDModel

oa::I32

—

inDFF

oa::I32

—

inNumExperts

oa::I32

—

inExpertsPerToken

oa::I32

—

inRmsEps

oa::F32

Default: 1e-5f

inNumSharedExperts

oa::I32

Default: 0

Returns

void

The declared return value.

const oa::Matrix & oa::Moe::lastSelectionMask()

No public source comment is attached to this declaration.

Returns

const oa::Matrix &

The declared return value.

oa::I32 oa::Moe::numExperts()

No public source comment is attached to this declaration.

Returns

oa::I32

The declared return value.

oa::I32 oa::Moe::numSharedExperts()

No public source comment is attached to this declaration.

Returns

oa::I32

The declared return value.

MoeRouteStats oa::Moe::routeStats()

── route telemetry (read after execute+Sync) ──────────────────────────────

Returns

MoeRouteStats

The declared return value.

oa::F32 oa::Moe::routingBias(
	oa::I32 inExpert
)

No public source comment is attached to this declaration.

Parameters

inExpert

oa::I32

—

Returns

oa::F32

The declared return value.

void oa::Moe::setAuxLossAlpha(
	oa::F32 inAlpha
)

── Optional differentiable balancing losses (opt-in, default off) ───────── Switch/GShard aux loss α·E·Σ_e f_e·p_e (f = hard load fraction, constant; P = mean router prob, differentiable) plus router z-loss β·mean(LSE²). Add auxLoss() to the task loss before backward. Both coefficients 0 ⇒ auxLoss() is a 0 scalar recorded on the tape with no gradient effect.

Parameters

inAlpha

oa::F32

—

Returns

void

The declared return value.

void oa::Moe::setBalanceRate(
	oa::F32 inGamma
)

── Aux-loss-free load balancing (DeepSeek-V3) ───────────────────────────── A per-expert bias added to the routing logits for the top-k SELECTION decision only (never into the gate magnitude, so it does not distort the weighted combine and produces no gradient). After each optimizer step, call updateRoutingBias() to nudge under-loaded experts up and over-loaded ones down. The update and next forward remain in the deferred GPU graph.

Parameters

inGamma

oa::F32

—

Returns

void

The declared return value.

void oa::Moe::setRouterZLossBeta(
	oa::F32 inBeta
)

No public source comment is attached to this declaration.

Parameters

inBeta

oa::F32

—

Returns

void

The declared return value.

void oa::Moe::setSparseExecution(
	bool inEnabled
)

No public source comment is attached to this declaration.

Parameters

inEnabled

bool

—

Returns

void

The declared return value.

bool oa::Moe::sparseExecution()

No public source comment is attached to this declaration.

Returns

bool

The declared return value.

void oa::Moe::updateRoutingBias()

No public source comment is attached to this declaration.

Returns

void

The declared return value.