Sparse MoE Routing

Task

Implement sparse_moe_routing(router_logits: list[list[float]], top_k: int, expert_capacity: int) -> tuple[list[list[int]], list[list[float]], list[int]]. Select and normalize each token's top-k experts, enforce per-expert capacity in deterministic token and route-rank order, and report accepted routes, weights, and final expert loads.

Requirements

• router_logits has shape tokens by experts; top_k is between zero and the expert count.
• Select experts by descending logit, breaking equal-logit ties by smaller expert index.
• Apply numerically stable softmax only across each token's selected top-k logits.
• Process assignments in token order and then selected route-rank order.
• Each expert accepts at most expert_capacity assignments.
• Represent a capacity-dropped assignment with expert index -1 and weight 0.
• Do not fall back to an unselected expert when a selected expert is full.
• Do not renormalize surviving weights after capacity drops.
• Return routed expert indices, routed weights, and accepted assignment counts for every expert.
• Do not mutate router_logits.

Example

router_logits = [
    [3.0, 2.0, 0.0],
    [4.0, 1.0, 3.0],
]

sparse_moe_routing(router_logits, top_k=2, expert_capacity=1)
# (
#   [[0, 1], [-1, 2]],
#   [[0.73105858, 0.26894142], [0.0, 0.26894142]],
#   [1, 1, 1],
# )