dockyard_rl.models.jax.moe.router¶
Token-choice top-K router for MoE (the gating network), Flax NNX.
JAX mirror of models/dtensor/moe/router.py: each token is scored against all
experts by a small bias-free gate, then routed to its top-K experts, with
optional node-limited (group-limited) routing for large-expert-count models. The
gate runs in float32 for load-balance stability (matching the torch CUDA
autocast override). Pure tensor logic — CPU-parity-validatable.
Shape suffixes: B=batch, L=seq, D=model dim, E=num experts, K=top-k.
Module Contents¶
Classes¶
Top-K token-choice gating. |
Data¶
API¶
- dockyard_rl.models.jax.moe.router.Array¶
None
- dockyard_rl.models.jax.moe.router.ScoreFunc¶
None
- class dockyard_rl.models.jax.moe.router.TokenChoiceTopKRouter(dim: int, num_experts: int, *, top_k: int = 1, score_func: dockyard_rl.models.jax.moe.router.ScoreFunc = 'sigmoid', route_norm: bool = False, route_scale: float = 1.0, num_expert_groups: Optional[int] = None, num_limited_groups: Optional[int] = None, rngs: flax.nnx.Rngs, param_dtype: jax.numpy.dtype = jnp.float32)¶
Bases:
flax.nnx.ModuleTop-K token-choice gating.
Args mirror the torch
TokenChoiceTopKRouter:score_funcselects the per-expert (sigmoid) vs normalized (softmax) gate activation;route_normrenormalizes the selected top-K weights to sum to 1;route_scalemultiplies the final weights;num_expert_groups/num_limited_groupsenable node-limited routing (group score = sum of the group’s top-2 expert scores).Initialization