MoE Router

The TopKRouter (src/xorl/models/layers/moe/router.py) is stateless — it owns no parameters. The gate linear (nn.Linear(hidden_size, num_experts, bias=False)) lives on MoEBlock so that the checkpoint path resolves to mlp.gate.weight, matching HuggingFace conventions.

---

Routing Algorithm

routing_weights = softmax(gate(hidden_states), dim=-1)   # [T, E]
routing_weights, selected_experts = topk(routing_weights, k=top_k, dim=-1)  # [T, k]
if norm_topk_prob:
    routing_weights /= routing_weights.sum(dim=-1, keepdim=True)  # renormalize to sum=1

Each token independently selects its top-K experts from the softmax distribution over all E experts. Gate scores are broadcast into the weighted sum after expert compute.

---

Router Parameters

Parameter	Type	Description
norm_topk_prob	bool	Renormalize the top-k gate scores to sum to 1. Qwen3-MoE uses true. With renormalization the output is a convex combination of expert outputs. Without it, scores are raw softmax values — which can be very small when E is large.
router_fp32	bool	Upcast gate linear and hidden states to float32 before computing logits. Recommended for numerical stability and alignment with SGLang inference.
freeze_router	bool	Freeze the gate linear weights during training (server training). Useful for fine-tuning where pre-trained routing should be preserved.
train_router	bool	When false, detaches routing weights before they flow into the expert weighted sum. The gate still trains via auxiliary losses on router_logits, but not through expert gradients. Stabilizes training when routing is already good.

model:
  router_fp32: true        # recommended for stability + SGLang alignment

freeze_router: true        # server training: freeze gate.weight entirely

train:
  no_decay_modules: [gate] # local training: exclude gate from weight decay

---

Routing Replay (R3)

Routing Replay (R3) solves a correctness problem that arises when combining EP with gradient checkpointing.

The problem: If recomputation changes the values at the routing boundary, top-K selection can pick different experts. That changes dispatch shapes and makes the recomputed execution inconsistent with the recorded forward.

The solution: R3 records routing decisions (expert indices, optionally gate weights) during the forward pass and replays them verbatim during backward recompute.

Lifecycle Stages

Stage	When	What happens
record	Forward pass	Router runs normally; top-K indices are stored to CPU-pinned buffer
replay_forward	Recompute in backward	Router result discarded; cached indices loaded instead
replay_backward	Backward through replayed forward	AllToAll shapes guaranteed identical to original forward

R3 Mode for RL (Pre-loaded Routing from SGLang)

R3 also powers rollout routing replay for RL training. When xorl-sglang exports expert choices, XoRL can load them so the route-conditioned trainer computation uses the same discrete assignments. This does not verify the router computation independently and does not prove equality of expert outputs, selected-token logprobs, gradients, or post-update weights.

The local forward/recompute replay lifecycle is enabled automatically on admitted EP + gradient-checkpointing paths. Rollout replay additionally requires routing data from the inference response and the corresponding server request path.

---

Source

File	Description
src/xorl/models/layers/moe/router.py	TopKRouter — softmax, top-K, optional renormalization
src/xorl/models/layers/moe/routing_replay.py	RoutingReplay — R3 record/replay for checkpoint determinism and RL routing consistency
src/xorl/models/layers/moe/moe_block.py	MoEBlock — gate linear ownership, router invocation, R3 integration