MoE: Dispatch and Weighted Combine
Medium · Solved in PyTorch · Machine learning coding practice
Problem
Implement the sparse execution step of an MoE layer. x contains token features of shape (T,D). expert_indices and expert_weights have shape (T,k); each row has distinct zero-based expert IDs and nonnegative weights summing to one. Dimensions are positive and 1 \leq k \leq E, where E is the number of supplied experts.
Each expert is a row-wise callable mapping a batch of shape (n,D) to the same shape without modifying its input. Call each selected expert once, with only its assigned rows of x in original token order. Skip experts with no assignments. Send unweighted features to the experts; apply the routing weights to their outputs. Keep zero-weight assignments in the dispatch and counts. No capacity limit or residual addition is needed.
Topics: mixture-of-experts, transformers, routing, pytorch, numpy
The full statement, worked examples, hints and the test suite are available once you sign in. You can then solve MoE: Dispatch and Weighted Combine in the browser and run it against the tests.